Update app.py

app.py

@@ -187,75 +187,68 @@ def predict_anime_aesthetic(img, model):
 #      Image Evaluation Tool        #
 #####################################
 
-class ImageEvaluationTool:
-    """Evaluation tool to process images through multiple aesthetic models and generate logs and HTML outputs."""
+class ModelManager:
+    """Manages model loading and processing requests using a queue."""
     def __init__(self):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print(f"Using device: {self.device}")
         print("Loading models... This may take some time.")
 
-        # Load models with progress logs
+        # Load models once during initialization
         print("Loading Aesthetic Shadow model...")
-        self.aesthetic_shadow = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=self.device)
+        self.aesthetic_shadow_model = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=self.device)
         print("Loading Waifu Scorer model...")
-        self.waifu_scorer = WaifuScorer(device=self.device, verbose=True)
+        self.waifu_scorer_model = WaifuScorer(device=self.device, verbose=True)
         print("Loading Aesthetic Predictor V2.5...")
-        self.aesthetic_predictor = load_aesthetic_predictor_v2_5()
+        self.aesthetic_predictor_model = load_aesthetic_predictor_v2_5()
         print("Loading Anime Aesthetic model...")
-        self.anime_aesthetic = load_anime_aesthetic_model()
+        self.anime_aesthetic_model = load_anime_aesthetic_model()
         print("All models loaded successfully!")
 
-        self.temp_dir = tempfile.mkdtemp()
-        self.results = []  # Store final results for sorting and display
         self.available_models = {
-            "aesthetic_shadow": {"name": "Aesthetic Shadow", "process": self._process_aesthetic_shadow},
-            "waifu_scorer": {"name": "Waifu Scorer", "process": self._process_waifu_scorer},
-            "aesthetic_predictor_v2_5": {"name": "Aesthetic V2.5", "process": self._process_aesthetic_predictor_v2_5},
-            "anime_aesthetic": {"name": "Anime Score", "process": self._process_anime_aesthetic},
+            "aesthetic_shadow": {"name": "Aesthetic Shadow", "process": self._process_aesthetic_shadow, "model": self.aesthetic_shadow_model},
+            "waifu_scorer": {"name": "Waifu Scorer", "process": self._process_waifu_scorer, "model": self.waifu_scorer_model},
+            "aesthetic_predictor_v2_5": {"name": "Aesthetic V2.5", "process": self._process_aesthetic_predictor_v2_5, "model": self.aesthetic_predictor_model},
+            "anime_aesthetic": {"name": "Anime Score", "process": self._process_anime_aesthetic, "model": self.anime_aesthetic_model},
         }
+        self.processing_queue: asyncio.Queue = asyncio.Queue()
+        self.worker_task = asyncio.create_task(self._worker())
+        self.temp_dir = tempfile.mkdtemp()
 
-    def image_to_base64(self, image: Image.Image) -> str:
-        """Convert PIL Image to base64 encoded JPEG string."""
-        buffered = BytesIO()
-        image.save(buffered, format="JPEG")
-        return base64.b64encode(buffered.getvalue()).decode('utf-8')
-
-    def auto_tune_batch_size(self, images: list) -> int:
-        """Automatically determine the optimal batch size for processing."""
-        batch_size = 1
-        max_batch = len(images)
-        test_image = images[0:1]
-        while batch_size <= max_batch:
-            try:
-                if "aesthetic_shadow" in self.available_models and self.available_models["aesthetic_shadow"]['selected']: # Check if model is available and selected
-                    _ = self.aesthetic_shadow(test_image * batch_size)
-                if "waifu_scorer" in self.available_models and self.available_models["waifu_scorer"]['selected']: # Check if model is available and selected
-                    _ = self.waifu_scorer(test_image * batch_size)
-                if "aesthetic_predictor_v2_5" in self.available_models and self.available_models["aesthetic_predictor_v2_5"]['selected']: # Check if model is available and selected
-                    _ = self.aesthetic_predictor.inference(test_image * batch_size)
-                batch_size *= 2
-                if batch_size > max_batch:
-                    break
-            except Exception:
+    async def _worker(self):
+        """Background worker to process image evaluation requests from the queue."""
+        while True:
+            request = await self.processing_queue.get()
+            if request is None: # Shutdown signal
+                self.processing_queue.task_done()
                 break
-        optimal = max(1, batch_size // 2)
-        if optimal > 64:
-            optimal = 64
-            print("Capped optimal batch size to 64")
-        print(f"Optimal batch size determined: {optimal}")
-        return optimal
-
-    async def process_images_evaluation_with_logs(self, file_paths: list, auto_batch: bool, manual_batch_size: int, selected_models):
-        """Asynchronously process images and yield updates with logs, HTML table, and progress bar."""
-        self.results = []
+            try:
+                results = await self._process_request(request)
+                request['results_future'].set_result(results) # Fulfill the future with results
+            except Exception as e:
+                request['results_future'].set_exception(e) # Set exception if processing fails
+            finally:
+                self.processing_queue.task_done()
+
+    async def submit_request(self, request_data):
+        """Submit a new image processing request to the queue."""
+        results_future = asyncio.Future() # Future to hold the results
+        request = {**request_data, 'results_future': results_future}
+        await self.processing_queue.put(request)
+        return await results_future # Wait for and return results
+
+    async def _process_request(self, request):
+        """Process a single image evaluation request."""
+        file_paths = request['file_paths']
+        auto_batch = request['auto_batch']
+        manual_batch_size = request['manual_batch_size']
+        selected_models = request['selected_models']
         log_events = []
         images = []
         file_names = []
+        final_results = []
 
-        # Update available models based on selection
-        for model_key in self.available_models:
-            self.available_models[model_key]['selected'] = model_key in selected_models
-
+        # Prepare images and file names
         total_files = len(file_paths)
         log_events.append(f"Starting to load {total_files} images...")
         for f in file_paths:
@@ -268,14 +261,9 @@ class ImageEvaluationTool:
 
         if not images:
             log_events.append("No valid images loaded.")
-            progress_percentage = 0 # Define progress_percentage here for no images case
-            progress_html = self._generate_progress_html(progress_percentage)
-            yield ("<p>No images loaded.</p>", "", self._format_logs(log_events), progress_html, manual_batch_size)
-            return
+            return [], log_events, 0, manual_batch_size
 
-        yield ("<p>Images loaded. Determining batch size...</p>", "", self._format_logs(log_events),
-               self._generate_progress_html(0), manual_batch_size)
-        await asyncio.sleep(0.1)
+        log_events.append("Images loaded. Determining batch size...")
 
         try:
             manual_batch_size = int(manual_batch_size) if manual_batch_size is not None else 1
@@ -285,9 +273,6 @@ class ImageEvaluationTool:
 
         optimal_batch = self.auto_tune_batch_size(images) if auto_batch else manual_batch_size
         log_events.append(f"Using batch size: {optimal_batch}")
-        yield ("<p>Processing images in batches...</p>", "", self._format_logs(log_events),
-               self._generate_progress_html(0), optimal_batch)
-        await asyncio.sleep(0.1)
 
         total_images = len(images)
         for i in range(0, total_images, optimal_batch):
@@ -298,36 +283,18 @@ class ImageEvaluationTool:
 
             batch_results = {}
 
-            # Aesthetic Shadow processing
-            if self.available_models['aesthetic_shadow']['selected']:
-                batch_results['aesthetic_shadow'] = await self._process_aesthetic_shadow(batch_images, log_events)
-            else:
-                batch_results['aesthetic_shadow'] = [None] * len(batch_images)
-
-            # Waifu Scorer processing
-            if self.available_models['waifu_scorer']['selected']:
-                batch_results['waifu_scorer'] = await self._process_waifu_scorer(batch_images, log_events)
-            else:
-                batch_results['waifu_scorer'] = [None] * len(batch_images)
-
-            # Aesthetic Predictor V2.5 processing
-            if self.available_models['aesthetic_predictor_v2_5']['selected']:
-                batch_results['aesthetic_predictor_v2_5'] = await self._process_aesthetic_predictor_v2_5(batch_images, log_events)
-            else:
-                batch_results['aesthetic_predictor_v2_5'] = [None] * len(batch_images)
-
-            # Anime Aesthetic processing (single image)
-            if self.available_models['anime_aesthetic']['selected']:
-                batch_results['anime_aesthetic'] = await self._process_anime_aesthetic(batch_images, log_events)
-            else:
-                batch_results['anime_aesthetic'] = [None] * len(batch_images)
+            # Process selected models
+            for model_key in selected_models:
+                if self.available_models[model_key]['selected']: # Ensure model is selected
+                    batch_results[model_key] = await self.available_models[model_key]['process'](self, batch_images, log_events)
+                else:
+                    batch_results[model_key] = [None] * len(batch_images)
 
-
-            # Combine results
+            # Combine results and create final results list
             for j in range(len(batch_images)):
                 scores_to_average = []
-                for model_key in self.available_models:
-                    if self.available_models[model_key]['selected']: # Only consider selected models
+                for model_key in selected_models:
+                     if self.available_models[model_key]['selected']: # Ensure model is selected
                         score = batch_results[model_key][j]
                         if score is not None:
                             scores_to_average.append(score)
@@ -340,28 +307,49 @@ class ImageEvaluationTool:
                     'img_data': self.image_to_base64(thumbnail),  # Keep this for the HTML display
                     'final_score': final_score,
                 }
-                for model_key in self.available_models: # Add model scores to result
+                for model_key in selected_models: # Add model scores to result
                     if self.available_models[model_key]['selected']:
                         result[model_key] = batch_results[model_key][j]
+                final_results.append(result)
 
-                self.results.append(result)
-                self.sort_results() # Sort results after adding new result
-                progress_percentage = min(100, ((i + len(batch_images)) / total_images) * 100) # Define progress_percentage here
-                yield (f"<p>Processed batch {batch_index}.</p>", self.generate_html_table(self.results, selected_models), # Update table immediately
-                       self._format_logs(log_events[-10:]), self._generate_progress_html(progress_percentage), optimal_batch)
-                await asyncio.sleep(0.1)
+        log_events.append("All images processed.")
+        return final_results, log_events, 100, optimal_batch
 
 
-        log_events.append("All images processed.")
-        self.sort_results() # Final sort after all images processed
-        html_table = self.generate_html_table(self.results, selected_models) # Pass selected models to final table generation
-        final_progress = self._generate_progress_html(100)
-        yield ("<p>All images processed.</p>", html_table,
-               self._format_logs(log_events[-10:]), final_progress, optimal_batch)
+    def image_to_base64(self, image: Image.Image) -> str:
+        """Convert PIL Image to base64 encoded JPEG string."""
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+    def auto_tune_batch_size(self, images: list) -> int:
+        """Automatically determine the optimal batch size for processing."""
+        batch_size = 1
+        max_batch = len(images)
+        test_image = images[0:1]
+        while batch_size <= max_batch:
+            try:
+                if "aesthetic_shadow" in self.available_models and self.available_models["aesthetic_shadow"]['selected']: # Check if model is available and selected
+                    _ = self.available_models["aesthetic_shadow"]['model'](test_image * batch_size)
+                if "waifu_scorer" in self.available_models and self.available_models["waifu_scorer"]['selected']: # Check if model is available and selected
+                    _ = self.available_models["waifu_scorer"]['model'](test_image * batch_size)
+                if "aesthetic_predictor_v2_5" in self.available_models and self.available_models["aesthetic_predictor_v2_5"]['selected']: # Check if model is available and selected
+                    _ = self.available_models["aesthetic_predictor_v2_5"]['model'].inference(test_image * batch_size)
+                batch_size *= 2
+                if batch_size > max_batch:
+                    break
+            except Exception:
+                break
+        optimal = max(1, batch_size // 2)
+        if optimal > 64:
+            optimal = 64
+            print("Capped optimal batch size to 64")
+        print(f"Optimal batch size determined: {optimal}")
+        return optimal
 
     async def _process_aesthetic_shadow(self, batch_images, log_events):
         try:
-            shadow_results = self.aesthetic_shadow(batch_images)
+            shadow_results = self.available_models["aesthetic_shadow"]['model'](batch_images)
             log_events.append("Aesthetic Shadow processed for batch.")
         except Exception as e:
             log_events.append(f"Error in Aesthetic Shadow: {e}")
@@ -379,7 +367,7 @@ class ImageEvaluationTool:
 
     async def _process_waifu_scorer(self, batch_images, log_events):
         try:
-            waifu_scores = self.waifu_scorer(batch_images)
+            waifu_scores = self.available_models["waifu_scorer"]['model'](batch_images)
             waifu_scores = [float(np.clip(s, 0.0, 10.0)) if s is not None else None for s in waifu_scores]
             log_events.append("Waifu Scorer processed for batch.")
         except Exception as e:
@@ -389,7 +377,7 @@ class ImageEvaluationTool:
 
     async def _process_aesthetic_predictor_v2_5(self, batch_images, log_events):
         try:
-            v2_5_scores = self.aesthetic_predictor.inference(batch_images)
+            v2_5_scores = self.available_models["aesthetic_predictor_v2_5"]['model'].inference(batch_images)
             v2_5_scores = [float(np.round(np.clip(s, 0.0, 10.0), 4)) if s is not None else None for s in v2_5_scores]
             log_events.append("Aesthetic Predictor V2.5 processed for batch.")
         except Exception as e:
@@ -401,7 +389,7 @@ class ImageEvaluationTool:
         anime_scores = []
         for j, img in enumerate(batch_images):
             try:
-                score = predict_anime_aesthetic(img, self.anime_aesthetic)
+                score = predict_anime_aesthetic(img, self.available_models["anime_aesthetic"]['model'])
                 anime_scores.append(float(np.clip(score * 10.0, 0.0, 10.0)))
                 log_events.append(f"Anime Aesthetic processed for image {j + 1}.")
             except Exception as e:
@@ -424,7 +412,7 @@ class ImageEvaluationTool:
         """Format log events into an HTML string."""
         return "<div style='max-height:300px; overflow-y:auto;'>" + "<br>".join(logs) + "</div>"
 
-    def sort_results(self, sort_by: str = "Final Score") -> list:
+    def sort_results(self, results, sort_by: str = "Final Score") -> list:
         """Sort results based on the specified column."""
         key_map = {
             "Final Score": "final_score",
@@ -436,8 +424,8 @@ class ImageEvaluationTool:
         }
         key = key_map.get(sort_by, "final_score")
         reverse = sort_by != "File Name"
-        self.results.sort(key=lambda r: r.get(key) if r.get(key) is not None else (-float('inf') if not reverse else float('inf')), reverse=reverse)
-        return self.results
+        results.sort(key=lambda r: r.get(key) if r.get(key) is not None else (-float('inf') if not reverse else float('inf')), reverse=reverse)
+        return results
 
     def generate_html_table(self, results: list, selected_models) -> str:
         """Generate an HTML table to display the evaluation results."""
@@ -503,17 +491,25 @@ class ImageEvaluationTool:
 
 
     def cleanup(self):
-        """Clean up temporary directories."""
+        """Clean up temporary directories and shutdown worker."""
         if os.path.exists(self.temp_dir):
             shutil.rmtree(self.temp_dir)
+        asyncio.run(self.shutdown()) # Shutdown worker gracefully
+
+    async def shutdown(self):
+        """Send shutdown signal to worker and wait for it to finish."""
+        await self.processing_queue.put(None) # Send shutdown signal
+        await self.worker_task # Wait for worker task to complete
+        await self.processing_queue.join() # Wait for queue to be empty
 
 
 #####################################
 #             Interface             #
 #####################################
 
+model_manager = ModelManager() # Initialize ModelManager once outside the interface function
+
 def create_interface():
-    evaluator = ImageEvaluationTool()
     sort_options = ["Final Score", "File Name", "Aesthetic Shadow", "Waifu Scorer", "Aesthetic V2.5", "Anime Score"]
     model_options = ["aesthetic_shadow", "waifu_scorer", "aesthetic_predictor_v2_5", "anime_aesthetic"]
 
@@ -556,12 +552,13 @@ def create_interface():
                 status_html = gr.HTML(label="Status")
                 output_html = gr.HTML(label="Evaluation Results")
                 download_file_output = gr.File() # Initialize gr.File component without filename
+                global_results_state = gr.State([]) # Initialize a global state to hold results
 
         # Function to convert results to CSV format, excluding 'img_data'.
-        def results_to_csv(selected_models):
+        def results_to_csv(results, selected_models): # Take results as input
             import csv
             import io
-            if not evaluator.results:
+            if not results:
                 return None # Return None when no results are available
             output = io.StringIO()
             fieldnames = ['file_name', 'final_score'] # Base fieldnames
@@ -571,7 +568,7 @@ def create_interface():
 
             writer = csv.DictWriter(output, fieldnames=fieldnames)
             writer.writeheader()
-            for res in evaluator.results:
+            for res in results:
                 row_dict = {'file_name': res['file_name'], 'final_score': res['final_score']} # Base data
                 for model_key in selected_models: # Add selected model scores
                     if model_key in selected_models: # Double check before accessing res[model_key]
@@ -583,33 +580,49 @@ def create_interface():
         def update_batch_size_interactivity(auto_batch):
             return gr.update(interactive=not auto_batch)
 
-        async def process_images_and_update(files, auto_batch, manual_batch, selected_models):
+        async def process_images_and_update(files, auto_batch, manual_batch, selected_models, current_results):
             file_paths = [f.name for f in files]
-            async for status, table, logs, progress, updated_batch in evaluator.process_images_evaluation_with_logs(file_paths, auto_batch, manual_batch, selected_models):
-                yield status, table, logs, progress, gr.update(value=updated_batch, interactive=not auto_batch)
 
-        def update_table_sort(sort_by_column, selected_models):
-            sorted_results = evaluator.sort_results(sort_by_column)
-            return evaluator.generate_html_table(sorted_results, selected_models)
+            # Prepare request data for the ModelManager
+            request_data = {
+                'file_paths': file_paths,
+                'auto_batch': auto_batch,
+                'manual_batch_size': manual_batch,
+                'selected_models': {model: {'selected': model in selected_models} for model in model_options} # Pass model selections
+            }
+            # Submit request and get results from ModelManager
+            results, logs, progress_percent, updated_batch = await model_manager.submit_request(request_data)
+
+            updated_results = current_results + results # Append new results to current results
+
+            html_table = model_manager.generate_html_table(updated_results, selected_models)
+            progress_html = model_manager._generate_progress_html(progress_percent)
+            log_html = model_manager._format_logs(logs[-10:])
 
-        def update_table_model_selection(selected_models):
+            return status_html, html_table, log_html, progress_html, gr.update(value=updated_batch, interactive=not auto_batch), updated_results
+
+
+        def update_table_sort(sort_by_column, selected_models, current_results):
+            sorted_results = model_manager.sort_results(current_results, sort_by_column)
+            return model_manager.generate_html_table(sorted_results, selected_models), sorted_results # Return sorted results
+
+        def update_table_model_selection(selected_models, current_results):
             # Recalculate final scores based on selected models
-            for result in evaluator.results:
+            for result in current_results:
                 scores_to_average = []
-                for model_key in evaluator.available_models:
-                    if model_key in selected_models and evaluator.available_models[model_key]['selected']: # consider only selected models from checkbox group and available_models
+                for model_key in model_options: # Use model_options here, not available_models from manager in UI context
+                    if model_key in selected_models and model_key in model_manager.available_models and model_manager.available_models[model_key]['selected']: # consider only selected models from checkbox group and available_models
                         score = result.get(model_key)
                         if score is not None:
                             scores_to_average.append(score)
                 final_score = float(np.clip(np.mean(scores_to_average), 0.0, 10.0)) if scores_to_average else None
                 result['final_score'] = final_score
 
-            sorted_results = evaluator.sort_results() # Keep sorting by Final Score when models change
-            return evaluator.generate_html_table(sorted_results, selected_models)
+            sorted_results = model_manager.sort_results(current_results, "Final Score") # Keep sorting by Final Score when models change
+            return model_manager.generate_html_table(sorted_results, selected_models), sorted_results
 
 
         def clear_results():
-            evaluator.results = []
             return (gr.update(value=""),
                     gr.update(value=""),
                     gr.update(value=""),
@@ -618,10 +631,11 @@ def create_interface():
                       <div style='width:0%;background-color:#4CAF50;padding:5px 0;text-align:center;'>0%</div>
                     </div>
                     """),
-                    gr.update(value=1))
+                    gr.update(value=1),
+                    []) # Clear results state
 
-        def download_results_csv_trigger(selected_models): # Changed function name to avoid conflict and clarify purpose
-            csv_content = results_to_csv(selected_models)
+        def download_results_csv_trigger(selected_models, current_results): # Changed function name to avoid conflict and clarify purpose
+            csv_content = results_to_csv(current_results, selected_models)
             if csv_content is None:
                 return None # Indicate no file to download
 
@@ -633,6 +647,10 @@ def create_interface():
             return temp_file_path # Return the path to the temporary file
 
 
+        # Set initial selection state for models in ModelManager (important!)
+        for model_key in model_options:
+            model_manager.available_models[model_key]['selected'] = True # Default to all selected initially
+
         auto_batch_checkbox.change(
             update_batch_size_interactivity,
             inputs=[auto_batch_checkbox],
@@ -641,30 +659,30 @@ def create_interface():
 
         process_btn.click(
             process_images_and_update,
-            inputs=[input_images, auto_batch_checkbox, batch_size_input, model_checkboxes],
-            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input]
+            inputs=[input_images, auto_batch_checkbox, batch_size_input, model_checkboxes, global_results_state],
+            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input, global_results_state]
         )
         sort_dropdown.change(
             update_table_sort,
-            inputs=[sort_dropdown, model_checkboxes],
-            outputs=[output_html]
+            inputs=[sort_dropdown, model_checkboxes, global_results_state],
+            outputs=[output_html, global_results_state]
         )
         model_checkboxes.change( # Added change event for model checkboxes
             update_table_model_selection,
-            inputs=[model_checkboxes],
-            outputs=[output_html]
+            inputs=[model_checkboxes, global_results_state],
+            outputs=[output_html, global_results_state]
         )
         clear_btn.click(
             clear_results,
             inputs=[],
-            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input]
+            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input, global_results_state]
         )
         download_csv.click(
             download_results_csv_trigger, # Call the trigger function
-            inputs=[model_checkboxes],
+            inputs=[model_checkboxes, global_results_state],
             outputs=[download_file_output] # Output is now the gr.File component
         )
-        demo.load(lambda: update_table_sort("Final Score", model_options), inputs=None, outputs=[output_html]) # Initial sort and table render
+        demo.load(lambda: update_table_sort("Final Score", model_options, []), inputs=None, outputs=[output_html, global_results_state]) # Initial sort and table render, pass empty initial results
         gr.Markdown("""
         ### Notes
         - Select models to use for evaluation using the checkboxes.