Update app.py

app.py

@@ -1,27 +1,32 @@
-import gradio as gr
-import torch
 import os
-import numpy as np
+import shutil
+import tempfile
+import base64
+import asyncio
+from io import BytesIO
+
 import cv2
+import numpy as np
+import torch
 import onnxruntime as rt
 from PIL import Image
+import gradio as gr
 from transformers import pipeline
 from huggingface_hub import hf_hub_download
-import pandas as pd
-import tempfile
-import shutil
-import base64
-from io import BytesIO
 
 # Import necessary function from aesthetic_predictor_v2_5
 from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip
 
+
+#####################################
+#         Model Definitions         #
+#####################################
+
 class MLP(torch.nn.Module):
-    def __init__(self, input_size, xcol='emb', ycol='avg_rating', batch_norm=True):
+    """A simple multi-layer perceptron for image feature regression."""
+    def __init__(self, input_size: int, batch_norm: bool = True):
         super().__init__()
         self.input_size = input_size
-        self.xcol = xcol
-        self.ycol = ycol
         self.layers = torch.nn.Sequential(
             torch.nn.Linear(self.input_size, 2048),
             torch.nn.ReLU(),
@@ -44,29 +49,37 @@ class MLP(torch.nn.Module):
             torch.nn.Linear(32, 1)
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.layers(x)
 
-class WaifuScorer(object):
-    def __init__(self, model_path=None, device='cuda', cache_dir=None, verbose=False):
+
+class WaifuScorer:
+    """WaifuScorer model that uses CLIP for feature extraction and a custom MLP for scoring."""
+    def __init__(self, model_path: str = None, device: str = 'cuda', cache_dir: str = None, verbose: bool = False):
         self.verbose = verbose
+        self.device = device
+        self.dtype = torch.float32
+        self.available = False
 
         try:
-            import clip
-
+            import clip  # local import to avoid dependency issues
+            # Set default model path if not provided
             if model_path is None:
                 model_path = "Eugeoter/waifu-scorer-v3/model.pth"
                 if self.verbose:
-                    print(f"model path not set, switch to default: `{model_path}`")
+                    print(f"Model path not provided. Using default: {model_path}")
 
+            # Download model if not found locally
             if not os.path.isfile(model_path):
-                split = model_path.split("/")
-                username, repo_id, model_name = split[-3], split[-2], split[-1]
+                username, repo_id, model_name = model_path.split("/")[-3:]
                 model_path = hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=cache_dir)
 
-            print(f"Loading WaifuScorer model from `{model_path}`")
+            if self.verbose:
+                print(f"Loading WaifuScorer model from: {model_path}")
 
+            # Initialize MLP model
             self.mlp = MLP(input_size=768)
+            # Load state dict
             if model_path.endswith(".safetensors"):
                 from safetensors.torch import load_file
                 state_dict = load_file(model_path)
@@ -74,42 +87,44 @@ class WaifuScorer(object):
                 state_dict = torch.load(model_path, map_location=device)
             self.mlp.load_state_dict(state_dict)
             self.mlp.to(device)
-
-            self.model2, self.preprocess = clip.load("ViT-L/14", device=device)
-            self.device = device
-            self.dtype = torch.float32
             self.mlp.eval()
+
+            # Load CLIP model for image preprocessing and feature extraction
+            self.clip_model, self.preprocess = clip.load("ViT-L/14", device=device)
             self.available = True
         except Exception as e:
             print(f"Unable to initialize WaifuScorer: {e}")
-            self.available = False
 
     @torch.no_grad()
     def __call__(self, images):
         if not self.available:
-            return [None] * (1 if not isinstance(images, list) else len(images))
-
+            return [None] * (len(images) if isinstance(images, list) else 1)
         if isinstance(images, Image.Image):
             images = [images]
         n = len(images)
+        # Ensure at least two images for CLIP model compatibility
         if n == 1:
-            images = images*2
+            images = images * 2
 
         image_tensors = [self.preprocess(img).unsqueeze(0) for img in images]
         image_batch = torch.cat(image_tensors).to(self.device)
-        image_features = self.model2.encode_image(image_batch)
-
-        l2 = image_features.norm(2, dim=-1, keepdim=True)
-        l2[l2 == 0] = 1
-        im_emb_arr = (image_features / l2).to(device=self.device, dtype=self.dtype)
-
-        predictions = self.mlp(im_emb_arr)
+        image_features = self.clip_model.encode_image(image_batch)
+        # Normalize features
+        norm = image_features.norm(2, dim=-1, keepdim=True)
+        norm[norm == 0] = 1
+        im_emb = (image_features / norm).to(device=self.device, dtype=self.dtype)
+        predictions = self.mlp(im_emb)
         scores = predictions.clamp(0, 10).cpu().numpy().reshape(-1).tolist()
-
         return scores[:n]
 
+
+#####################################
+#   Aesthetic Predictor Functions   #
+#####################################
+
 def load_aesthetic_predictor_v2_5():
-    class AestheticPredictorV2_5_Impl: # Renamed class to avoid confusion
+    """Load and return an instance of Aesthetic Predictor V2.5 with batch processing support."""
+    class AestheticPredictorV2_5_Impl:
         def __init__(self):
             print("Loading Aesthetic Predictor V2.5...")
             self.model, self.preprocessor = convert_v2_5_from_siglip(
@@ -119,516 +134,546 @@ def load_aesthetic_predictor_v2_5():
             if torch.cuda.is_available():
                 self.model = self.model.to(torch.bfloat16).cuda()
 
-        def inference(self, image: Image.Image) -> float:
-            # preprocess image
-            pixel_values = self.preprocessor(
-                images=image.convert("RGB"), return_tensors="pt"
-            ).pixel_values
-
-            if torch.cuda.is_available():
-                pixel_values = pixel_values.to(torch.bfloat16).cuda()
-
-            # predict aesthetic score
-            with torch.inference_mode():
-                score = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
+        def inference(self, image):
+            if isinstance(image, list):
+                images_rgb = [img.convert("RGB") for img in image]
+                pixel_values = self.preprocessor(images=images_rgb, return_tensors="pt").pixel_values
+                if torch.cuda.is_available():
+                    pixel_values = pixel_values.to(torch.bfloat16).cuda()
+                with torch.inference_mode():
+                    scores = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
+                if scores.ndim == 0:
+                    scores = np.array([scores])
+                return scores.tolist()
+            else:
+                pixel_values = self.preprocessor(images=image.convert("RGB"), return_tensors="pt").pixel_values
+                if torch.cuda.is_available():
+                    pixel_values = pixel_values.to(torch.bfloat16).cuda()
+                with torch.inference_mode():
+                    score = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
+                return score
 
-            return score
+    return AestheticPredictorV2_5_Impl()
 
-    return AestheticPredictorV2_5_Impl() # Return an instance of the implementation class
 
 def load_anime_aesthetic_model():
+    """Load and return the Anime Aesthetic ONNX model."""
     model_path = hf_hub_download(repo_id="skytnt/anime-aesthetic", filename="model.onnx")
-    model = rt.InferenceSession(model_path, providers=['CPUExecutionProvider'])
-    return model
+    return rt.InferenceSession(model_path, providers=['CPUExecutionProvider'])
+
 
 def predict_anime_aesthetic(img, model):
-    img = np.array(img).astype(np.float32) / 255
+    """Predict Anime Aesthetic score for a single image."""
+    img_np = np.array(img).astype(np.float32) / 255.0
     s = 768
-    h, w = img.shape[:-1]
-    h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
-    ph, pw = s - h, s - w
-    img_input = np.zeros([s, s, 3], dtype=np.float32)
-    img_input[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(img, (w, h))
-    img_input = np.transpose(img_input, (2, 0, 1))
-    img_input = img_input[np.newaxis, :]
-    pred = model.run(None, {"img": img_input})[0].item()
+    h, w = img_np.shape[:2]
+    if h > w:
+        new_h, new_w = s, int(s * w / h)
+    else:
+        new_h, new_w = int(s * h / w), s
+    resized = cv2.resize(img_np, (new_w, new_h))
+    # Center the resized image in a square canvas
+    canvas = np.zeros((s, s, 3), dtype=np.float32)
+    pad_h = (s - new_h) // 2
+    pad_w = (s - new_w) // 2
+    canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
+    # Prepare input for model
+    input_tensor = np.transpose(canvas, (2, 0, 1))[np.newaxis, :]
+    pred = model.run(None, {"img": input_tensor})[0].item()
     return pred
 
+
+#####################################
+#      Image Evaluation Tool        #
+#####################################
+
 class ImageEvaluationTool:
+    """Evaluation tool to process images through multiple aesthetic models and generate logs and HTML outputs."""
     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
         print("Loading Aesthetic Shadow model...")
         self.aesthetic_shadow = 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)
-
         print("Loading Aesthetic Predictor V2.5...")
-        self.aesthetic_predictor_v2_5 = load_aesthetic_predictor_v2_5()
-
+        self.aesthetic_predictor = load_aesthetic_predictor_v2_5()
         print("Loading Anime Aesthetic model...")
         self.anime_aesthetic = 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},
+        }
+
+    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 evaluate_image(self, image):
-        results = {}
+    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:
+                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 = []
+        log_events = []
+        images = []
+        file_names = []
+
+        # Update available models based on selection
+        for model_key in self.available_models:
+            self.available_models[model_key]['selected'] = model_key in selected_models
+
+        total_files = len(file_paths)
+        log_events.append(f"Starting to load {total_files} images...")
+        for f in file_paths:
+            try:
+                img = Image.open(f).convert("RGB")
+                images.append(img)
+                file_names.append(os.path.basename(f))
+            except Exception as e:
+                log_events.append(f"Error opening {f}: {e}")
 
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
+        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
 
-        try:
-            shadow_result = self.aesthetic_shadow(images=[image])[0]
-            hq_score = [p for p in shadow_result if p['label'] == 'hq'][0]['score']
-            # Scale aesthetic_shadow to 0-10 and clamp
-            aesthetic_shadow_score = np.clip(hq_score * 10.0, 0.0, 10.0)
-            results['aesthetic_shadow'] = aesthetic_shadow_score
-        except Exception as e:
-            print(f"Error in Aesthetic Shadow: {e}")
-            results['aesthetic_shadow'] = None
+        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)
 
         try:
-            waifu_score = self.waifu_scorer([image])[0]
-            # Clamp waifu_score
-            waifu_score_clamped = np.clip(waifu_score, 0.0, 10.0)
-            results['waifu_scorer'] = waifu_score_clamped
-        except Exception as e:
-            print(f"Error in Waifu Scorer: {e}")
-            results['waifu_scorer'] = None
+            manual_batch_size = int(manual_batch_size) if manual_batch_size is not None else 1
+        except ValueError:
+            manual_batch_size = 1
+            log_events.append("Invalid manual batch size. Defaulting to 1.")
+
+        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):
+            batch_images = images[i:i+optimal_batch]
+            batch_file_names = file_names[i:i+optimal_batch]
+            batch_index = i // optimal_batch + 1
+            log_events.append(f"Processing batch {batch_index}: images {i+1} to {min(i+optimal_batch, total_images)}")
+
+            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)
 
-        try:
-            v2_5_score = self.aesthetic_predictor_v2_5.inference(image)
-            # Clamp v2.5 score
-            v2_5_score_clamped = np.clip(v2_5_score, 0.0, 10.0)
-            results['aesthetic_predictor_v2_5'] = float(np.round(v2_5_score_clamped, 4)) # Keep 4 decimal places after clamping
-        except Exception as e:
-            print(f"Error in Aesthetic Predictor V2.5: {e}")
-            results['aesthetic_predictor_v2_5'] = None
+            # 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)
 
-        try:
-            img_array = np.array(image)
-            anime_score = predict_anime_aesthetic(img_array, self.anime_aesthetic)
-            # Scale Anime Score to 0-10 and clamp
-            anime_score_scaled = np.clip(anime_score * 10.0, 0.0, 10.0)
-            results['anime_aesthetic'] = anime_score_scaled
-        except Exception as e:
-            print(f"Error in Anime Aesthetic: {e}")
-            results['anime_aesthetic'] = None
+            # 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)
 
-        # Calculate Final Score (simple average of available scores)
-        valid_scores = [v for v in results.values() if v is not None]
-        if valid_scores:
-            final_score = np.mean(valid_scores)
-            results['final_score'] = np.clip(final_score, 0.0, 10.0) # Clamp final score too
-        else:
-            results['final_score'] = None
+            # 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)
 
-        return results
 
-    def image_to_base64(self, image):
-        buffered = BytesIO()
-        image.save(buffered, format="JPEG")
-        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+            # Combine results
+            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
+                        score = batch_results[model_key][j]
+                        if score is not None:
+                            scores_to_average.append(score)
 
-    def process_single_image(self, file_path):
+                final_score = float(np.clip(np.mean(scores_to_average), 0.0, 10.0)) if scores_to_average else None
+                thumbnail = batch_images[j].copy()
+                thumbnail.thumbnail((200, 200))
+                result = {
+                    'file_name': batch_file_names[j],
+                    '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
+                    if self.available_models[model_key]['selected']:
+                        result[model_key] = batch_results[model_key][j]
+
+                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.")
+        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)
+
+    async def _process_aesthetic_shadow(self, batch_images, log_events):
         try:
-            img = Image.open(file_path).convert("RGB")
-            eval_results = self.evaluate_image(img)
-            thumbnail = img.copy()
-            thumbnail.thumbnail((200, 200))
-            img_base64 = self.image_to_base64(thumbnail)
-            result = {
-                'file_name': os.path.basename(file_path),
-                'img_data': img_base64,
-                **eval_results
-            }
-            return result
+            shadow_results = self.aesthetic_shadow(batch_images)
+            log_events.append("Aesthetic Shadow processed for batch.")
         except Exception as e:
-            print(f"Error processing {file_path}: {e}")
-            return None
-
-    def process_images_evaluation(self, image_files): # Renamed and now for evaluation only
-        results = []
-
-        for i, file_path in enumerate(image_files):
+            log_events.append(f"Error in Aesthetic Shadow: {e}")
+            shadow_results = [None] * len(batch_images)
+        aesthetic_shadow_scores = []
+        for res in shadow_results:
             try:
-                img = Image.open(file_path).convert("RGB")
-                eval_results = self.evaluate_image(img)
-
-                thumbnail = img.copy()
-                thumbnail.thumbnail((200, 200))
-
-                img_base64 = self.image_to_base64(thumbnail)
+                hq_score = next(p for p in res if p['label'] == 'hq')['score']
+                score = float(np.clip(hq_score * 10.0, 0.0, 10.0))
+            except Exception:
+                score = None
+            aesthetic_shadow_scores.append(score)
+        log_events.append("Aesthetic Shadow scores computed for batch.")
+        return aesthetic_shadow_scores
+
+    async def _process_waifu_scorer(self, batch_images, log_events):
+        try:
+            waifu_scores = self.waifu_scorer(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:
+            log_events.append(f"Error in Waifu Scorer: {e}")
+            waifu_scores = [None] * len(batch_images)
+        return waifu_scores
 
-                result = {
-                    'file_name': os.path.basename(file_path),
-                    'img_data': img_base64,
-                    **eval_results
-                }
-                results.append(result)
+    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 = [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:
+            log_events.append(f"Error in Aesthetic Predictor V2.5: {e}")
+            v2_5_scores = [None] * len(batch_images)
+        return v2_5_scores
 
+    async def _process_anime_aesthetic(self, batch_images, log_events):
+        anime_scores = []
+        for j, img in enumerate(batch_images):
+            try:
+                score = predict_anime_aesthetic(img, self.anime_aesthetic)
+                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:
-                print(f"Error processing {file_path}: {e}")
-
-        return results
-
-    def sort_results(self, results, sort_by="Final Score"): # New function for sorting
-        def sort_key(res): # Define a sorting key function
-            sort_value = res.get(sort_by.lower().replace(" ", "_"), None) # Handle spaces and case
-            if sort_value is None: # Put N/A at the end
-                return -float('inf') if sort_by == "File Name" else float('inf') # File Name sort N/A at end alphabetically
-            return sort_value
-
-        results.sort(key=sort_key, reverse=sort_by != "File Name") # Sort results, reverse for score columns
-        return results
+                log_events.append(f"Error in Anime Aesthetic for image {j + 1}: {e}")
+                anime_scores.append(None)
+        return anime_scores
+
+
+    def _generate_progress_html(self, percentage: float) -> str:
+        """Generate HTML for a progress bar given a percentage."""
+        return f"""
+            <div style="width:100%;background-color:#ddd; border-radius:5px;">
+              <div style="width:{percentage:.1f}%; background-color:#4CAF50; text-align:center; padding:5px 0; border-radius:5px;">
+                {percentage:.1f}%
+              </div>
+            </div>
+        """
 
-    def generate_html_table(self, results):
-        html = """
+    def _format_logs(self, logs: list) -> str:
+        """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:
+        """Sort results based on the specified column."""
+        key_map = {
+            "Final Score": "final_score",
+            "File Name": "file_name",
+            "Aesthetic Shadow": "aesthetic_shadow",
+            "Waifu Scorer": "waifu_scorer",
+            "Aesthetic V2.5": "aesthetic_predictor_v2_5",
+            "Anime Score": "anime_aesthetic"
+        }
+        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
+
+    def generate_html_table(self, results: list, selected_models) -> str:
+        """Generate an HTML table to display the evaluation results."""
+        table_html = """
         <style>
-            .results-table {
-                width: 100%;
-                border-collapse: collapse;
-                margin: 20px 0;
-                font-family: Arial, sans-serif;
-                background-color: transparent;
-            }
-
-            .results-table th,
-            .results-table td {
-                color: #eee;
-                border: 1px solid #ddd;
-                padding: 8px;
-                text-align: center;
-                background-color: transparent;
-            }
-
-            .results-table th {
-                font-weight: bold;
-            }
-
-            .results-table tr:nth-child(even) {
-                background-color: transparent;
-            }
-
-            .results-table tr:hover {
-                background-color: rgba(255, 255, 255, 0.1);
-            }
-
-            .image-preview {
-                max-width: 150px;
-                max-height: 150px;
-                display: block;
-                margin: 0 auto;
-            }
-
-            .good-score {
-                color: #0f0;
-                font-weight: bold;
-            }
-            .bad-score {
-                color: #f00;
-                font-weight: bold;
-            }
-            .medium-score {
-                color: orange;
-                font-weight: bold;
-            }
+            .results-table { width: 100%; border-collapse: collapse; margin: 20px 0; font-family: Arial, sans-serif; }
+            .results-table th, .results-table td { color: #eee; border: 1px solid #ddd; padding: 8px; text-align: center; }
+            .results-table th { font-weight: bold; }
+            .results-table tr:nth-child(even) { background-color: transparent; }
+            .results-table tr:hover { background-color: rgba(255, 255, 255, 0.1); }
+            .image-preview { max-width: 150px; max-height: 150px; display: block; margin: 0 auto; }
+            .good-score { color: #0f0; font-weight: bold; }
+            .bad-score { color: #f00; font-weight: bold; }
+            .medium-score { color: orange; font-weight: bold; }
         </style>
-
         <table class="results-table">
             <thead>
                 <tr>
                     <th>Image</th>
                     <th>File Name</th>
-                    <th>Aesthetic Shadow</th>
-                    <th>Waifu Scorer</th>
-                    <th>Aesthetic V2.5</th>
-                    <th>Anime Score</th>
-                    <th>Final Score</th>
-                </tr>
-            </thead>
-            <tbody>
         """
+        visible_models = [] # Keep track of visible model columns
+        if "aesthetic_shadow" in selected_models:
+            table_html += "<th>Aesthetic Shadow</th>"
+            visible_models.append("aesthetic_shadow")
+        if "waifu_scorer" in selected_models:
+            table_html += "<th>Waifu Scorer</th>"
+            visible_models.append("waifu_scorer")
+        if "aesthetic_predictor_v2_5" in selected_models:
+            table_html += "<th>Aesthetic V2.5</th>"
+            visible_models.append("aesthetic_predictor_v2_5")
+        if "anime_aesthetic" in selected_models:
+            table_html += "<th>Anime Score</th>"
+            visible_models.append("anime_aesthetic")
+        table_html += "<th>Final Score</th>"
+        table_html += "</tr></thead><tbody>"
 
         for result in results:
-            html += "<tr>"
-            html += f'<td><img src="data:image/jpeg;base64,{result["img_data"]}" class="image-preview"></td>'
-            html += f'<td>{result["file_name"]}</td>'
-
-            score = result["aesthetic_shadow"]
-            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 4 else "bad-score"
-            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
-
-            score = result["waifu_scorer"]
-            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 5 else "bad-score"
-            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
-
-            score = result["aesthetic_predictor_v2_5"]
-            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 5 else "bad-score"
-            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
-
-            score = result["anime_aesthetic"]
-            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 5 else "bad-score"
-            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
-
-            score = result["final_score"]
-            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 5 else "bad-score"
-            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
-
-
-            html += "</tr>"
-
-        html += """
-            </tbody>
-        </table>
-        """
+            table_html += "<tr>"
+            table_html += f'<td><img src="data:image/jpeg;base64,{result["img_data"]}" class="image-preview"></td>'
+            table_html += f'<td>{result["file_name"]}</td>'
+            for model_key in visible_models: # Iterate through visible models only
+                score = result.get(model_key)
+                table_html += self._format_score_cell(score)
+
+            score = result.get("final_score")
+            table_html += self._format_score_cell(score)
+            table_html += "</tr>"
+        table_html += """</tbody></table>"""
+        return table_html
+
+    def _format_score_cell(self, score):
+        score_str = f"{score:.4f}" if isinstance(score, (int, float)) else "N/A"
+        score_class = ""
+        if isinstance(score, (int, float)):
+            if score >= 7:
+                score_class = "good-score"
+            elif score >= 5:
+                score_class = "medium-score"
+            else:
+                score_class = "bad-score"
+        return f'<td class="{score_class}">{score_str}</td>'
 
-        return html
 
     def cleanup(self):
+        """Clean up temporary directories."""
         if os.path.exists(self.temp_dir):
             shutil.rmtree(self.temp_dir)
 
-# Global variable to store evaluation results
-global_results = None
 
-def create_interface():
-    global global_results # Use the global variable
+#####################################
+#             Interface             #
+#####################################
 
+def create_interface():
     evaluator = ImageEvaluationTool()
-    sort_options = ["Final Score", "File Name", "Aesthetic Shadow", "Waifu Scorer", "Aesthetic V2.5", "Anime Score"] # Sort options
+    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"]
 
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("""
         # Comprehensive Image Evaluation Tool
 
-        Upload images to evaluate them using multiple aesthetic and quality prediction models:
-
-        - **Aesthetic Shadow**: Evaluates high-quality vs low-quality images (scaled to 0-10)
-        - **Waifu Scorer**: Rates anime/illustration quality from 0-10
-        - **Aesthetic Predictor V2.5**: General aesthetic quality prediction (clamped to 0-10)
-        - **Anime Aesthetic**: Specific model for anime style images (scaled and clamped to 0-10)
-        - **Final Score**: Average of available scores (clamped to 0-10)
-
-        Upload multiple images to get a comprehensive evaluation table. Scores are clamped to the range 0.0000 - 10.0000.
+        Upload images to evaluate them using multiple aesthetic and quality prediction models.
+
+        **New features:**
+        - **Dynamic Final Score:** Final score recalculates on model selection changes.
+        - **Model Selection:** Choose which models to use for evaluation.
+        - **Dynamic Table Updates:** Table updates automatically based on model selection.
+        - **Automatic Sorting:** Table is automatically sorted by 'Final Score'.
+        - **Detailed Logs:** See major processing events (limited to the last 10).
+        - **Progress Bar:** Visual indication of processing status.
+        - **Asynchronous Updates:** Streaming status and logs during processing.
+        - **Batch Size Controls:** Choose manual batch size or let the tool auto-detect it.
+        - **Download Results:** Export the evaluation results as CSV.
         """)
 
         with gr.Row():
             with gr.Column(scale=1):
-                input_images = gr.Files(label="Upload Images")
-                sort_dropdown = gr.Dropdown(sort_options, value="Final Score", label="Sort by") # Dropdown for sorting
+                input_images = gr.Files(label="Upload Images", file_count="multiple")
+                model_checkboxes = gr.CheckboxGroup(model_options, label="Select Models", value=model_options, info="Choose models for evaluation.")
+                auto_batch_checkbox = gr.Checkbox(label="Automatic Batch Size Detection", value=False, info="Enable to automatically determine the optimal batch size.")
+                batch_size_input = gr.Number(label="Batch Size", value=1, interactive=True, info="Manually specify the batch size if auto-detection is disabled.")
+                sort_dropdown = gr.Dropdown(sort_options, value="Final Score", label="Sort by", info="Select the column to sort results by.")
                 process_btn = gr.Button("Evaluate Images", variant="primary")
                 clear_btn = gr.Button("Clear Results")
+                download_csv = gr.Button("Download CSV", variant="secondary")
 
             with gr.Column(scale=2):
-                progress_html = gr.HTML(label="Progress") # Keep progress_html if you want to show initial progress
+                progress_bar = gr.HTML(label="Progress Bar", value="""
+                <div style='width:100%;background-color:#ddd;'>
+                  <div style='width:0%;background-color:#4CAF50;padding:5px 0;text-align:center;'>0%</div>
+                </div>
+                """)
+                log_window = gr.HTML(label="Detailed Logs", value="<div style='max-height:300px; overflow-y:auto;'>Logs will appear here...</div>")
+                status_html = gr.HTML(label="Status")
                 output_html = gr.HTML(label="Evaluation Results")
-
-        def process_images_and_update(files):
-            global global_results
+                download_file_output = gr.File() # Initialize gr.File component without filename
+
+        # Function to convert results to CSV format, excluding 'img_data'.
+        def results_to_csv(selected_models):
+            import csv
+            import io
+            if not evaluator.results:
+                return None # Return None when no results are available
+            output = io.StringIO()
+            fieldnames = ['file_name', 'final_score'] # Base fieldnames
+            for model_key in selected_models: # Add selected model names as fieldnames
+                if model_key in selected_models: # Double check if model_key is indeed in selected_models list
+                    fieldnames.append(model_key)
+
+            writer = csv.DictWriter(output, fieldnames=fieldnames)
+            writer.writeheader()
+            for res in evaluator.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]
+                        row_dict[model_key] = res.get(model_key, 'N/A') # Use get with default 'N/A' if model not in result (shouldn't happen but for safety)
+                writer.writerow(row_dict)
+            return output.getvalue()
+
+
+        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):
             file_paths = [f.name for f in files]
-            total_files = len(file_paths)
-            results = []
-            
-            if not file_paths:
-                global_results = []
-                yield "<p>No files uploaded.</p>", gr.update()
-                return
-
-            # Helper function to generate a styled progress bar HTML snippet.
-            def generate_progress_bar(percentage):
-                return f"""
-                <div style="background-color: #ddd; border-radius: 5px; width: 100%; margin: 10px 0;">
-                <div style="width: {percentage:.1f}%; background-color: #4CAF50; text-align: center; padding: 5px 0; border-radius: 5px;">
-                    {percentage:.1f}%
-                </div>
-                </div>
-                """
-            
-            total_models = 4  # Total number of model steps per image.
-            
-            for i, file_path in enumerate(file_paths):
-                file_name = os.path.basename(file_path)
-                try:
-                    img = Image.open(file_path).convert("RGB")
-                except Exception as e:
-                    yield f"<p>Error opening {file_name}: {e}</p>", gr.update()
-                    continue
-                
-                # Update overall progress before starting this image.
-                overall_percent = (i / total_files) * 100
-                overall_bar = generate_progress_bar(overall_percent)
-                progress_html = f"""
-                <p>{file_name}: Starting evaluation...</p>
-                <p>Overall Progress:</p>
-                {overall_bar}
-                """
-                yield progress_html, gr.update()
-                
-                # === Model Step 1: Aesthetic Shadow ===
-                model_index = 0
-                sub_percent = ((model_index + 1) / total_models) * 100
-                sub_bar = generate_progress_bar(sub_percent)
-                progress_html = f"""
-                <p>{file_name}: Evaluating <strong>Aesthetic Shadow</strong>...</p>
-                <p>Current Image Progress:</p>
-                {sub_bar}
-                <p>Overall Progress:</p>
-                {overall_bar}
-                """
-                yield progress_html, gr.update()
-                try:
-                    shadow_result = evaluator.aesthetic_shadow(images=[img])[0]
-                    hq_score = [p for p in shadow_result if p['label'] == 'hq'][0]['score']
-                    aesthetic_shadow_score = np.clip(hq_score * 10.0, 0.0, 10.0)
-                except Exception as e:
-                    print(f"Error in Aesthetic Shadow for {file_name}: {e}")
-                    aesthetic_shadow_score = None
-                yield f"<p>{file_name}: <strong>Aesthetic Shadow</strong> evaluation complete.</p>", gr.update()
-                
-                # === Model Step 2: Waifu Scorer ===
-                model_index = 1
-                sub_percent = ((model_index + 1) / total_models) * 100
-                sub_bar = generate_progress_bar(sub_percent)
-                progress_html = f"""
-                <p>{file_name}: Evaluating <strong>Waifu Scorer</strong>...</p>
-                <p>Current Image Progress:</p>
-                {sub_bar}
-                <p>Overall Progress:</p>
-                {overall_bar}
-                """
-                yield progress_html, gr.update()
-                try:
-                    waifu_score = evaluator.waifu_scorer([img])[0]
-                    waifu_score = np.clip(waifu_score, 0.0, 10.0)
-                except Exception as e:
-                    print(f"Error in Waifu Scorer for {file_name}: {e}")
-                    waifu_score = None
-                yield f"<p>{file_name}: <strong>Waifu Scorer</strong> evaluation complete.</p>", gr.update()
-                
-                # === Model Step 3: Aesthetic Predictor V2.5 ===
-                model_index = 2
-                sub_percent = ((model_index + 1) / total_models) * 100
-                sub_bar = generate_progress_bar(sub_percent)
-                progress_html = f"""
-                <p>{file_name}: Evaluating <strong>Aesthetic Predictor V2.5</strong>...</p>
-                <p>Current Image Progress:</p>
-                {sub_bar}
-                <p>Overall Progress:</p>
-                {overall_bar}
-                """
-                yield progress_html, gr.update()
-                try:
-                    v2_5_score = evaluator.aesthetic_predictor_v2_5.inference(img)
-                    v2_5_score = float(np.round(np.clip(v2_5_score, 0.0, 10.0), 4))
-                except Exception as e:
-                    print(f"Error in Aesthetic Predictor V2.5 for {file_name}: {e}")
-                    v2_5_score = None
-                yield f"<p>{file_name}: <strong>Aesthetic Predictor V2.5</strong> evaluation complete.</p>", gr.update()
-                
-                # === Model Step 4: Anime Aesthetic ===
-                model_index = 3
-                sub_percent = ((model_index + 1) / total_models) * 100
-                sub_bar = generate_progress_bar(sub_percent)
-                progress_html = f"""
-                <p>{file_name}: Evaluating <strong>Anime Aesthetic</strong>...</p>
-                <p>Current Image Progress:</p>
-                {sub_bar}
-                <p>Overall Progress:</p>
-                {overall_bar}
-                """
-                yield progress_html, gr.update()
-                try:
-                    img_array = np.array(img)
-                    anime_score = predict_anime_aesthetic(img_array, evaluator.anime_aesthetic)
-                    anime_score = np.clip(anime_score * 10.0, 0.0, 10.0)
-                except Exception as e:
-                    print(f"Error in Anime Aesthetic for {file_name}: {e}")
-                    anime_score = None
-                yield f"<p>{file_name}: <strong>Anime Aesthetic</strong> evaluation complete.</p>", gr.update()
-                
-                # === Final Score Calculation and Results Collection ===
-                valid_scores = [v for v in [aesthetic_shadow_score, waifu_score, v2_5_score, anime_score] if v is not None]
-                final_score = np.clip(np.mean(valid_scores), 0.0, 10.0) if valid_scores else None
-                
-                # Create a thumbnail and store the evaluation results.
-                thumbnail = img.copy()
-                thumbnail.thumbnail((200, 200))
-                img_base64 = evaluator.image_to_base64(thumbnail)
-                result = {
-                    'file_name': file_name,
-                    'img_data': img_base64,
-                    'aesthetic_shadow': aesthetic_shadow_score,
-                    'waifu_scorer': waifu_score,
-                    'aesthetic_predictor_v2_5': v2_5_score,
-                    'anime_aesthetic': anime_score,
-                    'final_score': final_score
-                }
-                results.append(result)
-                
-                # Update overall progress for the processed file.
-                overall_percent = ((i + 1) / total_files) * 100
-                overall_bar = generate_progress_bar(overall_percent)
-                progress_html = f"""
-                <p>{file_name}: Evaluation complete. ({i + 1}/{total_files} images processed.)</p>
-                <p>Overall Progress:</p>
-                {overall_bar}
-                """
-                # Sort the results by Final Score and update the table.
-                sorted_results = evaluator.sort_results(results.copy(), sort_by="Final Score")
-                html_table = evaluator.generate_html_table(sorted_results)
-                yield progress_html, html_table
-            
-            # Sort final results by Final Score.
-            global_results = evaluator.sort_results(results, sort_by="Final Score")
-            yield "<p>All images processed.</p>", evaluator.generate_html_table(global_results)
-
-        def update_table_sort(sort_by_column): # New function for sorting update
-            global global_results
-            if global_results is None:
-                return "No images evaluated yet." # Or handle case when no images are evaluated
-            sorted_results = evaluator.sort_results(global_results, sort_by=sort_by_column)
-            html_table = evaluator.generate_html_table(sorted_results)
-            return html_table
+            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 clear_results():
-            global global_results
-            global_results = None # Clear stored results
-            return gr.update(value=""), gr.update(value="")
+        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)
 
+        def update_table_model_selection(selected_models):
+            # Recalculate final scores based on selected models
+            for result in evaluator.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
+                        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)
+
+
+        def clear_results():
+            evaluator.results = []
+            return (gr.update(value=""),
+                    gr.update(value=""),
+                    gr.update(value=""),
+                    gr.update(value="""
+                    <div style='width:100%;background-color:#ddd;'>
+                      <div style='width:0%;background-color:#4CAF50;padding:5px 0;text-align:center;'>0%</div>
+                    </div>
+                    """),
+                    gr.update(value=1))
+
+        def download_results_csv_trigger(selected_models): # Changed function name to avoid conflict and clarify purpose
+            csv_content = results_to_csv(selected_models)
+            if csv_content is None:
+                return None # Indicate no file to download
+
+            # Create a temporary file to save the CSV data
+            with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as tmp_file:
+                tmp_file.write(csv_content.encode())
+                temp_file_path = tmp_file.name # Get the path to the temporary file
+
+            return temp_file_path # Return the path to the temporary file
+
+
+        auto_batch_checkbox.change(
+            update_batch_size_interactivity,
+            inputs=[auto_batch_checkbox],
+            outputs=[batch_size_input]
+        )
 
         process_btn.click(
             process_images_and_update,
-            inputs=[input_images],
-            outputs=[progress_html, output_html]
+            inputs=[input_images, auto_batch_checkbox, batch_size_input, model_checkboxes],
+            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input]
         )
-        sort_dropdown.change( # Only update table on sort change
+        sort_dropdown.change(
             update_table_sort,
-            inputs=[sort_dropdown],
-            outputs=[output_html] # Only update output_html
+            inputs=[sort_dropdown, model_checkboxes],
+            outputs=[output_html]
+        )
+        model_checkboxes.change( # Added change event for model checkboxes
+            update_table_model_selection,
+            inputs=[model_checkboxes],
+            outputs=[output_html]
         )
         clear_btn.click(
             clear_results,
             inputs=[],
-            outputs=[progress_html, output_html]
+            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input]
         )
-
-        demo.load(lambda: None, inputs=None, outputs=None)
-
+        download_csv.click(
+            download_results_csv_trigger, # Call the trigger function
+            inputs=[model_checkboxes],
+            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
         gr.Markdown("""
         ### Notes
-        - The evaluation may take some time depending on the number and size of images
-        - For best results, use high-quality images
-        - Scores are color-coded: green for good (>=7), orange for medium (>=5), and red for poor scores (<5, or <4 for Aesthetic Shadow)
-        - Some models may fail for certain image types, shown as "N/A" in the results
-        - "Final Score" is a simple average of available model scores.
-        - Table is sortable by clicking the dropdown above the "Evaluate Images" button. Default sort is by "Final Score". Sorting happens instantly without re-evaluating images.
+        - Select models to use for evaluation using the checkboxes.
+        - The 'Final Score' recalculates dynamically when models are selected/deselected.
+        - The table updates automatically when models are selected/deselected and is always sorted by 'Final Score'.
+        - The log window displays the most recent 10 events.
+        - The progress bar shows overall processing status.
+        - When 'Automatic Batch Size Detection' is enabled, the batch size field becomes disabled.
+        - Use the download button to export your evaluation results as CSV.
         """)
 
     return demo