Update app.py

app.py

@@ -1,73 +1,76 @@
 import os
-import shutil
+import io
 import tempfile
-import asyncio
-from io import BytesIO, StringIO
-import csv
-from pathlib import Path
-import logging
+import shutil # Kept for potential future use, but not actively used for now.
 
 import cv2
 import numpy as np
+import pandas as pd
 import torch
 import onnxruntime as rt
 from PIL import Image
 import gradio as gr
-from transformers import pipeline, AutoProcessor, AutoModelForImageClassification
+from transformers import pipeline
 from huggingface_hub import hf_hub_download
 
-# Configure basic logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# --- Dependency: aesthetic_predictor_v2_5.py ---
-# This file should exist in the same directory or be in PYTHONPATH.
-# For demonstration, a stub is provided. Replace with actual implementation.
-# aesthetic_predictor_v2_5.py STUB START
-# (Normally this would be in its own file: aesthetic_predictor_v2_5.py)
-def convert_v2_5_from_siglip(repo_id="unum-cloud/siglip-base-patch16-224-aesthetic-v2.5", low_cpu_mem_usage=True, trust_remote_code=True):
-    logger.info(f"Loading model and preprocessor from Hugging Face Hub: {repo_id}")
-    try:
-        # Attempt to load actual models if available and network permits
-        processor = AutoProcessor.from_pretrained(repo_id, low_cpu_mem_usage=low_cpu_mem_usage, trust_remote_code=trust_remote_code)
-        model = AutoModelForImageClassification.from_pretrained(repo_id, low_cpu_mem_usage=low_cpu_mem_usage, trust_remote_code=trust_remote_code)
-        logger.info("Successfully loaded model and preprocessor from Hugging Face Hub.")
-    except Exception as e:
-        logger.warning(f"Failed to load from {repo_id} due to: {e}. Using fallback mock objects.")
-        # Fallback to simpler mock objects if HF download fails or for offline use
-        class MockProcessor:
-            def __call__(self, images, return_tensors="pt"):
-                if isinstance(images, list):
-                    num_images = len(images)
-                    return {"pixel_values": torch.randn(num_images, 3, 224, 224)}
-                else:
-                    return {"pixel_values": torch.randn(1, 3, 224, 224)}
-        class MockModel:
-            def __init__(self): self._parameters = {"dummy": torch.nn.Parameter(torch.empty(0))}
-            def __call__(self, pixel_values):
-                bs = pixel_values.shape[0]
+# Assuming aesthetic_predictor_v2_5.py is in the same directory or Python path.
+# If it's not available, the AestheticPredictorV25 model will fail to load.
+# For this example, a mock will be used if the real import fails.
+try:
+    from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip
+except ImportError:
+    print("Warning: aesthetic_predictor_v2_5.py not found. Using a mock for AestheticPredictorV25.")
+    def convert_v2_5_from_siglip(low_cpu_mem_usage=True, trust_remote_code=True):
+        # This is a mock.
+        mock_model_output = torch.randn(1, 1) # Represents logits for a single image
+        
+        class MockModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.dummy_param = torch.nn.Parameter(torch.empty(0)) # To have a device property
+
+            def forward(self, pixel_values):
+                # Return something that has .logits
+                # Batch size from pixel_values
+                batch_size = pixel_values.size(0)
+                # Create a namedtuple or simple class to mimic HuggingFace output object with .logits
                 class Output:
-                    def __init__(self, logits_val): self.logits = logits_val
-                return Output(logits_val=torch.rand(bs, 1) * 10) # Simulate scores 0-10
-            def to(self, *args, **kwargs): return self
-            def cuda(self, *args, **kwargs): return self
-            def bfloat16(self, *args, **kwargs): return self
-        processor = MockProcessor()
-        model = MockModel()
-        logger.info("Using fallback mock model and preprocessor for Aesthetic Predictor V2.5.")
-    return model, processor
-# aesthetic_predictor_v2_5.py STUB END
-
-
-#####################################
-#         Model Definitions         #
-#####################################
+                    pass
+                output = Output()
+                output.logits = torch.randn(batch_size, 1).to(self.dummy_param.device)
+                return output
+            
+            def to(self, device_or_dtype): # Simplified .to()
+                if isinstance(device_or_dtype, torch.dtype):
+                     # In a real scenario, handle dtype conversion
+                    pass
+                elif isinstance(device_or_dtype, str) or isinstance(device_or_dtype, torch.device):
+                    self.dummy_param = torch.nn.Parameter(torch.empty(0, device=device_or_dtype)) # Move dummy param to device
+                return self
+
+            def cuda(self): # Mock .cuda()
+                return self.to(torch.device('cuda'))
+
+
+        mock_model_instance = MockModel()
+        
+        # Mock preprocessor that returns a dict with "pixel_values"
+        mock_preprocessor = lambda images, return_tensors: {"pixel_values": torch.randn(len(images) if isinstance(images, list) else 1, 3, 224, 224)}
+        return mock_model_instance, mock_preprocessor
+
+# --- Configuration ---
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+DTYPE_WAIFU = torch.float32 # Specific dtype for WaifuScorer's MLP
+CACHE_DIR = None # Set to a path string to use a specific Hugging Face cache directory, e.g., "./hf_cache"
+
+# --- Model Definitions ---
 
 class MLP(torch.nn.Module):
+    """Custom MLP for WaifuScorer."""
     def __init__(self, input_size: int, batch_norm: bool = True):
         super().__init__()
         self.input_size = input_size
-        layers = [
+        self.layers = torch.nn.Sequential(
             torch.nn.Linear(self.input_size, 2048), torch.nn.ReLU(),
             torch.nn.BatchNorm1d(2048) if batch_norm else torch.nn.Identity(), torch.nn.Dropout(0.3),
             torch.nn.Linear(2048, 512), torch.nn.ReLU(),
@@ -78,916 +81,461 @@ class MLP(torch.nn.Module):
             torch.nn.BatchNorm1d(128) if batch_norm else torch.nn.Identity(), torch.nn.Dropout(0.1),
             torch.nn.Linear(128, 32), torch.nn.ReLU(),
             torch.nn.Linear(32, 1)
-        ]
-        self.layers = torch.nn.Sequential(*layers)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.layers(x)
-
+        )
+    def forward(self, x: torch.Tensor) -> torch.Tensor: return self.layers(x)
 
-class WaifuScorer:
-    def __init__(self, model_path: str = None, device: str = 'cuda', cache_dir: str = None, verbose: bool = False):
-        self.verbose = verbose
+class BaseImageScorer:
+    """Abstract base class for image scorers."""
+    def __init__(self, model_key: str, model_display_name: str, device: str = DEVICE, verbose: bool = False):
+        self.model_key = model_key
+        self.model_display_name = model_display_name
         self.device = device
-        self.dtype = torch.float32
-        self.available = False
-        self.clip_model = None
-        self.preprocess = None
-        self.mlp = None
-
-        try:
-            import clip  # Dynamically import clip
-            if model_path is None:
-                model_path = "Eugeoter/waifu-scorer-v3/model.pth"
-                if self.verbose: logger.info(f"WaifuScorer model path not provided. Using default: {model_path}")
+        self.verbose = verbose
+        self.model = None
+        self.preprocessor = None
+        self._load_model()
 
-            if not Path(model_path).is_file():
-                try:
-                    # Assuming model_path like "user/repo/file.pth" for hf_hub_download
-                    parts = model_path.split("/")
-                    if len(parts) >= 3:
-                        repo_id_parts = parts[:-1]
-                        filename = parts[-1]
-                        repo_id_str = "/".join(repo_id_parts)
-                        model_path_resolved = hf_hub_download(repo_id=repo_id_str, filename=filename, cache_dir=cache_dir)
-                    else: # try as repo_id and assume model.pth or common name
-                         model_path_resolved = hf_hub_download(repo_id=model_path, filename="model.pth", cache_dir=cache_dir) # fallback filename
-                except Exception as e:
-                    logger.error(f"Failed to download WaifuScorer model from HF Hub ({model_path}): {e}")
-                    # Try a more specific default if the generic one failed
-                    logger.info("Attempting to download specific WaifuScorer model Eugeoter/waifu-scorer-v3/model.pth")
-                    model_path_resolved = hf_hub_download("Eugeoter/waifu-scorer-v3", "model.pth", cache_dir=cache_dir)
-                model_path = model_path_resolved
+    def _load_model(self): raise NotImplementedError
+    def predict(self, images: list[Image.Image]) -> list[float | None]: raise NotImplementedError
 
+    def __call__(self, images: list[Image.Image]) -> list[float | None]:
+        if not self.model:
+            if self.verbose: print(f"{self.model_display_name} model not loaded.")
+            return [None] * len(images)
+        
+        rgb_images = [img.convert("RGB") if img.mode != "RGB" else img for img in images]
+        return self.predict(rgb_images)
 
-            if self.verbose: logger.info(f"Loading WaifuScorer model from: {model_path}")
+class WaifuScorerModel(BaseImageScorer):
+    def _load_model(self):
+        try:
+            import clip
+            model_hf_path = "Eugeoter/waifu-scorer-v3/model.pth" # Default path
+            
+            repo_id, filename = os.path.split(model_hf_path)
+            actual_model_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=CACHE_DIR)
+            if self.verbose: print(f"Loading WaifuScorer MLP from: {actual_model_path}")
 
-            self.mlp = MLP(input_size=768)
-            if str(model_path).endswith(".safetensors"):
+            self.mlp = MLP(input_size=768) # ViT-L/14 embedding size
+            if actual_model_path.endswith(".safetensors"):
                 from safetensors.torch import load_file
-                state_dict = load_file(model_path, device=device)
+                state_dict = load_file(actual_model_path, device=self.device)
             else:
-                state_dict = torch.load(model_path, map_location=device)
-            
-            # Adjust keys if necessary (e.g. if saved from DataParallel)
-            if any(key.startswith("module.") for key in state_dict.keys()):
-                state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
-
+                state_dict = torch.load(actual_model_path, map_location=self.device)
             self.mlp.load_state_dict(state_dict)
-            self.mlp.to(device=self.device, dtype=self.dtype)
-            self.mlp.eval()
+            self.mlp.to(self.device).eval()
 
-            self.clip_model, self.preprocess = clip.load("ViT-L/14", device=self.device)
-            self.available = True
-            logger.info("WaifuScorer initialized successfully.")
+            if self.verbose: print("Loading CLIP model ViT-L/14 for WaifuScorer.")
+            self.model, self.preprocessor = clip.load("ViT-L/14", device=self.device) # self.model is CLIP model
+            self.model.eval()
         except ImportError:
-            logger.error("OpenAI CLIP library not found. WaifuScorer will be unavailable. Please install with 'pip install openai-clip'")
+            if self.verbose: print("CLIP library not found. WaifuScorer will not be available.")
         except Exception as e:
-            logger.error(f"Unable to initialize WaifuScorer: {e}")
+            if self.verbose: print(f"Error loading WaifuScorer ({self.model_display_name}): {e}")
 
     @torch.no_grad()
-    def __call__(self, images: list[Image.Image]) -> list[float | None]:
-        if not self.available:
-            return [None] * len(images)
+    def predict(self, images: list[Image.Image]) -> list[float | None]:
+        if not self.model or not self.mlp: return [None] * len(images)
         
-        if not images:
-            return []
-
         original_n = len(images)
-        # Handle single image case for CLIP if it has issues with batch_size=1 (some versions might)
-        processed_images = images if len(images) > 1 else images * 2
+        processed_images = list(images)
+        if original_n == 1: processed_images.append(images[0]) # Duplicate for single image batch
         
         try:
-            image_tensors = [self.preprocess(img).unsqueeze(0) for img in processed_images]
-            image_batch = torch.cat(image_tensors).to(self.device)
-            image_features = self.clip_model.encode_image(image_batch)
-            
+            image_tensors = torch.cat([self.preprocessor(img).unsqueeze(0) for img in processed_images]).to(self.device)
+            image_features = self.model.encode_image(image_tensors)
             norm = image_features.norm(p=2, dim=-1, keepdim=True)
-            norm = torch.where(norm == 0, torch.tensor(1.0, device=norm.device, dtype=norm.dtype), norm) # Avoid division by zero
-            im_emb = (image_features / norm).to(device=self.device, dtype=self.dtype)
+            norm[norm == 0] = 1e-6 # Avoid division by zero, use small epsilon
+            im_emb = (image_features / norm).to(device=self.device, dtype=DTYPE_WAIFU)
             
             predictions = self.mlp(im_emb)
-            scores = predictions.clamp(0, 10).cpu().numpy().reshape(-1).tolist()
+            scores = predictions.clamp(0, 10).cpu().numpy().flatten().tolist()
             return scores[:original_n]
         except Exception as e:
-            logger.error(f"Error during WaifuScorer prediction: {e}")
+            if self.verbose: print(f"Error during {self.model_display_name} prediction: {e}")
             return [None] * original_n
 
-
-class AestheticPredictorV2_5_Wrapper:
-    def __init__(self, device: str):
-        self.device = device
-        self.model, self.preprocessor = convert_v2_5_from_siglip(
-            low_cpu_mem_usage=True, trust_remote_code=True
-        )
-        if self.device == 'cuda' and torch.cuda.is_available():
-            self.model = self.model.to(torch.bfloat16).cuda()
-        logger.info("Aesthetic Predictor V2.5 Wrapper initialized.")
+class AestheticPredictorV25(BaseImageScorer):
+    def _load_model(self):
+        try:
+            if self.verbose: print(f"Loading {self.model_display_name}...")
+            self.model, self.preprocessor = convert_v2_5_from_siglip(low_cpu_mem_usage=True, trust_remote_code=True)
+            # Model's .to() method should handle dtype (e.g. bfloat16) and device.
+            self.model = self.model.to(self.device) 
+            if self.device == 'cuda' and torch.cuda.is_available() and hasattr(self.model, 'to'): # some models might need explicit dtype
+                 self.model = self.model.to(torch.bfloat16)
+            self.model.eval()
+        except Exception as e:
+            if self.verbose: print(f"Error loading {self.model_display_name}: {e}")
 
     @torch.no_grad()
-    def inference(self, images: list[Image.Image]) -> list[float | None]:
-        if not images:
-            return []
+    def predict(self, images: list[Image.Image]) -> list[float | None]:
+        if not self.model or not self.preprocessor: return [None] * len(images)
         try:
-            images_rgb = [img.convert("RGB") for img in images]
-            pixel_values = self.preprocessor(images=images_rgb, return_tensors="pt").pixel_values
-            if self.device == 'cuda' and torch.cuda.is_available():
-                pixel_values = pixel_values.to(torch.bfloat16).cuda()
+            inputs = self.preprocessor(images=images, return_tensors="pt")
+            pixel_values = inputs["pixel_values"].to(self.model.dummy_param.device if hasattr(self.model, 'dummy_param') else self.device) # Use model's device
+            if self.device == 'cuda' and torch.cuda.is_available() and pixel_values.dtype != torch.bfloat16 : # Match dtype if model changed it
+                 pixel_values = pixel_values.to(torch.bfloat16)
+
+            output = self.model(pixel_values)
+            scores_tensor = output.logits if hasattr(output, 'logits') else output
+            scores = scores_tensor.squeeze().float().cpu().numpy()
             
-            scores_tensor = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
-            if scores_tensor.ndim == 0: # Single image result
-                scores = [scores_tensor.item()]
-            else:
-                scores = scores_tensor.tolist()
-            return [round(max(0.0, min(s, 10.0)), 4) for s in scores] # Clip and round
+            scores_list = [float(np.round(np.clip(s, 0.0, 10.0), 4)) for s in np.atleast_1d(scores)]
+            return scores_list
         except Exception as e:
-            logger.error(f"Error during Aesthetic Predictor V2.5 inference: {e}")
+            if self.verbose: print(f"Error during {self.model_display_name} prediction: {e}")
             return [None] * len(images)
 
-def load_anime_aesthetic_onnx_model(cache_dir: str = None) -> rt.InferenceSession | None:
-    try:
-        model_path = hf_hub_download(repo_id="skytnt/anime-aesthetic", filename="model.onnx", cache_dir=cache_dir)
-        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if torch.cuda.is_available() else ['CPUExecutionProvider']
-        session = rt.InferenceSession(model_path, providers=providers)
-        logger.info(f"Anime Aesthetic ONNX model loaded with providers: {session.get_providers()}")
-        return session
-    except Exception as e:
-        logger.error(f"Failed to load Anime Aesthetic ONNX model: {e}")
-        return None
-
-def preprocess_anime_aesthetic_batch(images_pil: list[Image.Image], target_size: int = 768) -> np.ndarray | None:
-    if not images_pil:
-        return None
-    batch_canvases = []
-    try:
-        for img_pil in images_pil:
-            img_np = np.array(img_pil.convert("RGB")).astype(np.float32) / 255.0
-            h, w = img_np.shape[:2]
-            if h > w:
-                new_h, new_w = target_size, int(target_size * w / h)
-            else:
-                new_h, new_w = int(target_size * h / w), target_size
-            
-            resized = cv2.resize(img_np, (new_w, new_h), interpolation=cv2.INTER_AREA)
-            canvas = np.zeros((target_size, target_size, 3), dtype=np.float32)
-            pad_h = (target_size - new_h) // 2
-            pad_w = (target_size - new_w) // 2
-            canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
-            batch_canvases.append(canvas)
-        
-        input_tensor_batch = np.array(batch_canvases, dtype=np.float32) # (N, H, W, C)
-        input_tensor_batch = np.transpose(input_tensor_batch, (0, 3, 1, 2)) # (N, C, H, W)
-        return input_tensor_batch
-    except Exception as e:
-        logger.error(f"Error during Anime Aesthetic preprocessing: {e}")
-        return None
-
-#####################################
-#      Image Evaluation Tool        #
-#####################################
-
-class ModelManager:
-    def __init__(self, cache_dir: str = None):
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        logger.info(f"Using device: {self.device}")
-        self.cache_dir = cache_dir
-        self.models = {}
-        self.model_configs = {}
-        self._load_all_models()
-
-        self.processing_queue: asyncio.Queue = asyncio.Queue()
-        self.worker_task = None
-        self._temp_files_to_clean = [] # For CSV files
-
-    def _load_all_models(self):
-        logger.info("Loading Aesthetic Shadow model...")
-        try:
-            self.models["aesthetic_shadow"] = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=0 if self.device == 'cuda' else -1)
-            self.model_configs["aesthetic_shadow"] = {"name": "Aesthetic Shadow", "process_func": self._process_aesthetic_shadow}
-            logger.info("Aesthetic Shadow model loaded.")
-        except Exception as e:
-            logger.error(f"Failed to load Aesthetic Shadow model: {e}")
-
-        logger.info("Loading Waifu Scorer model...")
-        try:
-            ws = WaifuScorer(device=self.device, cache_dir=self.cache_dir, verbose=True)
-            if ws.available:
-                self.models["waifu_scorer"] = ws
-                self.model_configs["waifu_scorer"] = {"name": "Waifu Scorer", "process_func": self._process_waifu_scorer}
-                logger.info("Waifu Scorer model loaded.")
-            else:
-                logger.warning("Waifu Scorer model is not available.")
-        except Exception as e:
-            logger.error(f"Failed to load Waifu Scorer model: {e}")
-
-        logger.info("Loading Aesthetic Predictor V2.5...")
-        try:
-            ap_v25 = AestheticPredictorV2_5_Wrapper(device=self.device)
-            self.models["aesthetic_predictor_v2_5"] = ap_v25
-            self.model_configs["aesthetic_predictor_v2_5"] = {"name": "Aesthetic V2.5", "process_func": self._process_aesthetic_predictor_v2_5}
-            logger.info("Aesthetic Predictor V2.5 loaded.")
-        except Exception as e:
-            logger.error(f"Failed to load Aesthetic Predictor V2.5: {e}")
-
-        logger.info("Loading Anime Aesthetic model...")
+class AnimeAestheticONNX(BaseImageScorer):
+    def _load_model(self):
         try:
-            aa_model = load_anime_aesthetic_onnx_model(cache_dir=self.cache_dir)
-            if aa_model:
-                self.models["anime_aesthetic"] = aa_model
-                self.model_configs["anime_aesthetic"] = {"name": "Anime Score", "process_func": self._process_anime_aesthetic}
-                logger.info("Anime Aesthetic model loaded.")
-            else:
-                logger.warning("Anime Aesthetic ONNX model failed to load and will be unavailable.")
+            if self.verbose: print(f"Loading {self.model_display_name} (ONNX)...")
+            model_path = hf_hub_download(repo_id="skytnt/anime-aesthetic", filename="model.onnx", cache_dir=CACHE_DIR)
+            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if self.device == 'cuda' else ['CPUExecutionProvider']
+            valid_providers = [p for p in providers if p in rt.get_available_providers()] or ['CPUExecutionProvider']
+            self.model = rt.InferenceSession(model_path, providers=valid_providers)
+            if self.verbose: print(f"{self.model_display_name} loaded with providers: {self.model.get_providers()}")
         except Exception as e:
-            logger.error(f"Failed to load Anime Aesthetic model: {e}")
-        
-        logger.info(f"Available models for processing: {list(self.model_configs.keys())}")
-
-
-    async def start_worker_if_not_running(self):
-        if self.worker_task is None or self.worker_task.done():
-            self.worker_task = asyncio.create_task(self._worker())
-            logger.info("Async worker started.")
-
-    async def _worker(self):
-        while True:
-            request = await self.processing_queue.get()
-            if request is None:  # Shutdown signal
-                self.processing_queue.task_done()
-                logger.info("Async worker received shutdown signal.")
-                break
-            
-            future = request.get('future')
-            try:
-                if request['type'] == 'run_evaluation_generator':
-                    # The generator itself is created here and returned via future
-                    # The Gradio callback will iterate over it
-                    gen = self.run_evaluation_generator(**request['params'])
-                    future.set_result(gen)
-                else:
-                    logger.warning(f"Unknown request type in worker: {request.get('type')}")
-                    if future: future.set_exception(ValueError("Unknown request type"))
-            except Exception as e:
-                logger.error(f"Error in worker processing request: {e}", exc_info=True)
-                if future: future.set_exception(e)
-            finally:
-                self.processing_queue.task_done()
-
-    async def submit_evaluation_request(self, file_paths, auto_batch, manual_batch_size, selected_model_keys):
-        await self.start_worker_if_not_running()
-        future = asyncio.Future()
-        request_item = {
-            'type': 'run_evaluation_generator',
-            'params': {
-                'file_paths': file_paths,
-                'auto_batch': auto_batch,
-                'manual_batch_size': manual_batch_size,
-                'selected_model_keys': selected_model_keys,
-            },
-            'future': future
-        }
-        await self.processing_queue.put(request_item)
-        return await future # Future resolves to the async generator
-
-    def auto_tune_batch_size(self, images: list[Image.Image], selected_model_keys: list[str]) -> int:
-        if not images or not selected_model_keys:
-            return 1
+            if self.verbose: print(f"Error loading {self.model_display_name}: {e}")
+
+    def _preprocess_image(self, img: Image.Image) -> np.ndarray:
+        img_np = np.array(img).astype(np.float32) / 255.0
+        s = 768
+        h, w = img_np.shape[:2]
+        r = min(s/h, s/w)
+        new_h, new_w = int(h*r), int(w*r)
         
-        batch_size = 1
-        max_possible_batch = len(images)
-        test_image_pil = [images[0].copy()] # A list containing one PIL image, copy to avoid issues with transforms
-
-        logger.info(f"Auto-tuning batch size with selected models: {selected_model_keys}")
+        resized = cv2.resize(img_np, (new_w, new_h), interpolation=cv2.INTER_AREA if r < 1 else cv2.INTER_LANCZOS4)
         
-        optimal_batch_size = 1
-        while batch_size <= max_possible_batch:
-            current_test_batch = test_image_pil * batch_size
-            try:
-                logger.debug(f"Testing batch size: {batch_size}")
-                if "aesthetic_shadow" in selected_model_keys and "aesthetic_shadow" in self.models:
-                    _ = self.models["aesthetic_shadow"](current_test_batch, batch_size=batch_size)
-                if "waifu_scorer" in selected_model_keys and "waifu_scorer" in self.models:
-                    _ = self.models["waifu_scorer"](current_test_batch)
-                if "aesthetic_predictor_v2_5" in selected_model_keys and "aesthetic_predictor_v2_5" in self.models:
-                    _ = self.models["aesthetic_predictor_v2_5"].inference(current_test_batch)
-                if "anime_aesthetic" in selected_model_keys and "anime_aesthetic" in self.models:
-                    processed_input = preprocess_anime_aesthetic_batch(current_test_batch)
-                    if processed_input is None: raise ValueError("Anime aesthetic preprocessing failed for test batch")
-                    _ = self.models["anime_aesthetic"].run(None, {"img": processed_input})
-                
-                optimal_batch_size = batch_size # This batch size worked
-                if batch_size * 2 > max_possible_batch : # If next step exceeds max, current is best fit
-                    if max_possible_batch > batch_size: # Check if we can exactly fit max_possible_batch
-                         # Test max_possible_batch one last time if it's > current batch_size and < batch_size*2
-                        pass # Current optimal_batch_size is good, or we can check max_possible_batch specifically
-                    break 
-                batch_size *= 2
-
-            except Exception as e: # Typically torch.cuda.OutOfMemoryError or similar
-                logger.warning(f"Auto-tune failed at batch size {batch_size} for at least one model: {e}")
-                break # Current optimal_batch_size is the largest that worked before this failure
-        
-        # Cap the batch size for very large numbers of images / powerful GPUs
-        final_optimal_batch = min(optimal_batch_size, max_possible_batch, 64)
-        logger.info(f"Optimal batch size determined: {final_optimal_batch}")
-        return max(1, final_optimal_batch)
-
+        canvas = np.zeros((s, s, 3), dtype=np.float32) # Fill with black
+        pad_h, pad_w = (s - new_h) // 2, (s - new_w) // 2
+        canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
+        return np.transpose(canvas, (2, 0, 1))[np.newaxis, :]
 
-    async def run_evaluation_generator(self, file_paths: list[str], auto_batch: bool, 
-                                 manual_batch_size: int, selected_model_keys: list[str]):
-        
-        log_messages = []
-        def _log(msg):
-            log_messages.append(msg)
-            logger.info(msg)
-
-        _log("Starting image evaluation...")
-        yield {"type": "log_update", "messages": log_messages[-20:]} # Show last 20 logs
-        yield {"type": "progress", "value": 0.0, "desc": "Initiating..."}
-
-        images_pil = []
-        file_names = []
-        for f_path_str in file_paths:
+    def predict(self, images: list[Image.Image]) -> list[float | None]:
+        if not self.model: return [None] * len(images)
+        scores = []
+        for img in images:
             try:
-                p = Path(f_path_str)
-                img = Image.open(p).convert("RGB")
-                images_pil.append(img)
-                file_names.append(p.name)
-                _log(f"Loaded image: {p.name}")
+                input_tensor = self._preprocess_image(img)
+                pred = self.model.run(None, {"img": input_tensor})[0].item()
+                scores.append(float(np.clip(pred * 10.0, 0.0, 10.0)))
             except Exception as e:
-                _log(f"Error opening {f_path_str}: {e}")
-        
-        yield {"type": "log_update", "messages": log_messages[-20:]}
-
-        if not images_pil:
-            _log("No valid images loaded. Aborting.")
-            yield {"type": "log_update", "messages": log_messages[-20:]}
-            yield {"type": "progress", "value": 1.0, "desc": "No images loaded"}
-            yield {"type": "final_results_state", "data": []} # ensure state is empty
-            return
-
-        actual_batch_size = 1
-        if auto_batch:
-            _log("Auto-tuning batch size...")
-            yield {"type": "log_update", "messages": log_messages[-20:]}
-            yield {"type": "progress", "value": 0.05, "desc": "Auto-tuning batch size..."}
-            actual_batch_size = self.auto_tune_batch_size(images_pil, selected_model_keys)
-            _log(f"Auto-detected batch size: {actual_batch_size}")
-        else:
-            actual_batch_size = int(manual_batch_size) if manual_batch_size > 0 else 1
-            _log(f"Using manual batch size: {actual_batch_size}")
-        
-        yield {"type": "batch_size_update", "value": actual_batch_size}
-        yield {"type": "log_update", "messages": log_messages[-20:]}
-
-        all_results_for_state = [] # Full data for gr.State
-        dataframe_rows_so_far = [] # Data for gr.DataFrame (PIL images, strings, numbers)
-
-        total_images = len(images_pil)
-        processed_count = 0
-
-        for i in range(0, total_images, actual_batch_size):
-            batch_images_pil = images_pil[i:i+actual_batch_size]
-            batch_file_names = file_names[i:i+actual_batch_size]
-            num_in_batch = len(batch_images_pil)
-            _log(f"Processing batch {i//actual_batch_size + 1}/{ (total_images + actual_batch_size -1) // actual_batch_size }: images {i+1} to {i+num_in_batch}")
-            yield {"type": "log_update", "messages": log_messages[-20:]}
-
-            batch_model_scores = {key: [None] * num_in_batch for key in self.model_configs.keys()}
-
-            for model_key in selected_model_keys:
-                if model_key in self.models and model_key in self.model_configs:
-                    _log(f"  Running {self.model_configs[model_key]['name']} for batch...")
-                    yield {"type": "log_update", "messages": log_messages[-20:]}
-                    try:
-                        scores = await self.model_configs[model_key]['process_func'](batch_images_pil)
-                        batch_model_scores[model_key] = scores
-                        _log(f"  {self.model_configs[model_key]['name']} scores: {scores}")
-                    except Exception as e:
-                        _log(f"  Error processing batch with {self.model_configs[model_key]['name']}: {e}")
-                        batch_model_scores[model_key] = [None] * num_in_batch # Ensure it's list of Nones
-                    yield {"type": "log_update", "messages": log_messages[-20:]}
-            
-            # Assemble results for this batch
-            current_batch_df_rows = []
-            for j in range(num_in_batch):
-                result_item_state = {'file_name': batch_file_names[j]} # For gr.State
-                
-                # For DataFrame: [PIL.Image, filename, score1, score2, ..., final_score]
-                thumbnail = batch_images_pil[j].copy()
-                thumbnail.thumbnail((150, 150)) # Create thumbnail
-                result_item_df_row = [thumbnail, batch_file_names[j]]
-
-
-                current_image_scores = []
-                for model_key in self.model_configs.keys(): # Iterate in defined order for consistency
-                    score = batch_model_scores[model_key][j]
-                    result_item_state[model_key] = score # For gr.State
-                    if model_key in selected_model_keys: # Only add to DF if selected
-                        result_item_df_row.append(f"{score:.4f}" if isinstance(score, (float, int)) else "N/A")
-                        if isinstance(score, (float, int)) and model_key in selected_model_keys:
-                             current_image_scores.append(score)
-                
-                final_score = None
-                if current_image_scores:
-                    final_score_val = float(np.mean([s for s in current_image_scores if s is not None]))
-                    final_score = float(np.clip(final_score_val, 0.0, 10.0))
-                
-                result_item_state['final_score'] = final_score
-                result_item_df_row.append(f"{final_score:.4f}" if final_score is not None else "N/A")
-                
-                all_results_for_state.append(result_item_state)
-                current_batch_df_rows.append(result_item_df_row)
-            
-            dataframe_rows_so_far.extend(current_batch_df_rows)
-            
-            processed_count += num_in_batch
-            progress_value = processed_count / total_images
-            yield {"type": "partial_results_df_rows", "data": dataframe_rows_so_far, "selected_model_keys": selected_model_keys}
-            yield {"type": "progress", "value": progress_value, "desc": f"Processed {processed_count}/{total_images}"}
-
-        _log("All images processed.")
-        yield {"type": "log_update", "messages": log_messages[-20:]}
-        yield {"type": "progress", "value": 1.0, "desc": "Completed!"}
-        yield {"type": "final_results_state", "data": all_results_for_state}
+                if self.verbose: print(f"Error predicting with {self.model_display_name} for one image: {e}")
+                scores.append(None)
+        return scores
 
+class AestheticShadowPipeline(BaseImageScorer):
+    def _load_model(self):
+        try:
+            if self.verbose: print(f"Loading {self.model_display_name} pipeline...")
+            pipeline_device = 0 if self.device == 'cuda' else -1 
+            self.model = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=pipeline_device)
+        except Exception as e:
+            if self.verbose: print(f"Error loading {self.model_display_name}: {e}")
 
-    async def _process_aesthetic_shadow(self, batch_images: list[Image.Image]) -> list[float | None]:
-        model = self.models.get("aesthetic_shadow")
-        if not model: return [None] * len(batch_images)
-        results = model(batch_images, batch_size=len(batch_images)) # Assuming pipeline can take batch_size hint
+    def predict(self, images: list[Image.Image]) -> list[float | None]:
+        if not self.model: return [None] * len(images)
         scores = []
-        for res_group in results: # Results might be List[List[Dict]] or List[Dict]
-            # Handle both single image and batch results from pipeline
-            current_res_list = res_group if isinstance(res_group, list) else [res_group]
-            try:
-                hq_score_item = next(p for p in current_res_list if p['label'] == 'hq')
-                score = float(np.clip(hq_score_item['score'] * 10.0, 0.0, 10.0))
-            except (StopIteration, KeyError, TypeError):
-                score = None
-            scores.append(score)
-        return scores
-
-    async def _process_waifu_scorer(self, batch_images: list[Image.Image]) -> list[float | None]:
-        model = self.models.get("waifu_scorer")
-        if not model: return [None] * len(batch_images)
-        raw_scores = model(batch_images)
-        return [float(np.clip(s, 0.0, 10.0)) if s is not None else None for s in raw_scores]
-
-    async def _process_aesthetic_predictor_v2_5(self, batch_images: list[Image.Image]) -> list[float | None]:
-        model = self.models.get("aesthetic_predictor_v2_5")
-        if not model: return [None] * len(batch_images)
-        # Already returns clipped & rounded scores or Nones
-        return model.inference(batch_images)
-
-    async def _process_anime_aesthetic(self, batch_images: list[Image.Image]) -> list[float | None]:
-        model = self.models.get("anime_aesthetic")
-        if not model: return [None] * len(batch_images)
-        
-        input_data = preprocess_anime_aesthetic_batch(batch_images)
-        if input_data is None:
-            return [None] * len(batch_images)
-            
         try:
-            preds = model.run(None, {"img": input_data})[0] # Assuming output is (N, 1) or (N,)
-            scores = [float(np.clip(p.item() * 10.0, 0.0, 10.0)) for p in preds]
-            return scores
+            pipeline_results = self.model(images, top_k=None) # Assuming pipeline handles batching
+            
+            # Ensure consistent output structure from pipeline (List[List[Dict]] vs List[Dict])
+            if images and pipeline_results and not isinstance(pipeline_results[0], list):
+                pipeline_results = [pipeline_results] 
+
+            for res_set in pipeline_results:
+                try:
+                    hq_score_dict = next(p for p in res_set if p['label'] == 'hq')
+                    scores.append(float(np.clip(hq_score_dict['score'] * 10.0, 0.0, 10.0)))
+                except (StopIteration, TypeError, KeyError): scores.append(None)
         except Exception as e:
-            logger.error(f"Error during Anime Aesthetic ONNX prediction: {e}")
-            return [None] * len(batch_images)
+            if self.verbose: print(f"Error during {self.model_display_name} prediction: {e}")
+            return [None] * len(images) # All None if batch fails
+        return scores
 
-    def add_temp_file_for_cleanup(self, file_path: str):
-        self._temp_files_to_clean.append(file_path)
+# --- Model Management ---
+MODEL_REGISTRY = {
+    "aesthetic_shadow": {"class": AestheticShadowPipeline, "name": "Aesthetic Shadow"},
+    "waifu_scorer": {"class": WaifuScorerModel, "name": "Waifu Scorer"},
+    "aesthetic_predictor_v2_5": {"class": AestheticPredictorV25, "name": "Aesthetic V2.5"},
+    "anime_aesthetic": {"class": AnimeAestheticONNX, "name": "Anime Score"},
+}
+LOADED_MODELS = {} # Populated at startup
+
+def initialize_models(verbose_loading=False):
+    print(f"Using device: {DEVICE}")
+    print("Initializing models...")
+    for key, config in MODEL_REGISTRY.items():
+        LOADED_MODELS[key] = config["class"](key, config['name'], device=DEVICE, verbose=verbose_loading)
+    print("Model initialization complete.")
+
+# --- Core Logic ---
+@torch.no_grad()
+def auto_tune_batch_size(images: list[Image.Image], selected_model_keys: list[str], 
+                         initial_bs: int = 1, max_bs_limit: int = 64, verbose: bool = False) -> int:
+    if not images or not selected_model_keys: return initial_bs
+    if verbose: print("Auto-tuning batch size...")
+    
+    test_image = images[0]
+    active_models = [LOADED_MODELS[key] for key in selected_model_keys if key in LOADED_MODELS and LOADED_MODELS[key].model]
+    if not active_models: return initial_bs
 
-    async def shutdown_worker(self):
-        if self.worker_task and not self.worker_task.done():
-            logger.info("Attempting to shutdown worker...")
-            await self.processing_queue.put(None) # Send shutdown signal
+    bs = initial_bs
+    optimal_bs = initial_bs
+    while bs <= len(images) and bs <= max_bs_limit:
+        try:
+            batch_test_images = [test_image] * bs
+            for model in active_models:
+                if verbose: print(f"  Testing {model.model_display_name} with batch size {bs}")
+                model.predict(batch_test_images)
+            if DEVICE == 'cuda': torch.cuda.empty_cache()
+            
+            optimal_bs = bs
+            if bs == max_bs_limit: break
+            bs = min(bs * 2, max_bs_limit) # Try next power of 2 or max_bs_limit
+        except Exception as e: # Typically OOM or other runtime errors
+            if verbose: print(f"  Failed at batch size {bs} ({type(e).__name__}). Optimal so far: {optimal_bs}. Error: {str(e)[:100]}")
+            break 
+    if verbose: print(f"Auto-tuned batch size: {optimal_bs}")
+    return max(1, optimal_bs)
+
+async def evaluate_images_core(
+    pil_images: list[Image.Image], file_names: list[str], 
+    selected_model_keys: list[str], batch_size: int, 
+    progress_tracker: gr.Progress
+) -> tuple[pd.DataFrame, list[str]]:
+    
+    logs = []
+    num_images = len(pil_images)
+    if num_images == 0: return pd.DataFrame(), ["No images to process."]
+
+    # Initialize results_data: list of dicts, one per image
+    results_data = [{'File Name': fn, 'Thumbnail': img.copy().resize((150,150)), 'Final Score': np.nan} 
+                    for fn, img in zip(file_names, pil_images)]
+    for r_dict in results_data: # Initialize all model score columns to NaN
+        for cfg in MODEL_REGISTRY.values(): r_dict[cfg['name']] = np.nan
+
+    progress_tracker(0, desc="Starting evaluation...")
+    total_models_to_run = len(selected_model_keys)
+
+    for model_idx, model_key in enumerate(selected_model_keys):
+        model = LOADED_MODELS.get(model_key)
+        if not model or not model.model:
+            logs.append(f"Skipping {MODEL_REGISTRY[model_key]['name']} (not loaded).")
+            continue
+
+        model_name = model.model_display_name
+        logs.append(f"Processing with {model_name}...")
+        
+        current_img_offset = 0
+        for batch_start_idx in range(0, num_images, batch_size):
+            # Progress: (current_model_idx + fraction_of_current_model_done) / total_models_to_run
+            model_progress_fraction = (batch_start_idx / num_images)
+            overall_progress = (model_idx + model_progress_fraction) / total_models_to_run
+            progress_tracker(overall_progress, desc=f"{model_name} (Batch {batch_start_idx//batch_size + 1})")
+
+            batch_images = pil_images[batch_start_idx : batch_start_idx + batch_size]
             try:
-                await asyncio.wait_for(self.worker_task, timeout=5.0)
-                logger.info("Worker task finished.")
-            except asyncio.TimeoutError:
-                logger.warning("Worker task did not finish in time. Cancelling...")
-                self.worker_task.cancel()
+                scores = model(batch_images) # Use __call__
+                for i, score in enumerate(scores):
+                    results_data[current_img_offset + i][model_name] = score if score is not None else np.nan
             except Exception as e:
-                logger.error(f"Exception during worker shutdown: {e}")
-        await self.processing_queue.join() # Wait for queue to be fully processed
-        logger.info("Processing queue joined.")
-        self.worker_task = None
-
-
-    def cleanup(self):
-        logger.info("Running cleanup...")
-        # Shut down asyncio worker
-        if self.worker_task:
-             # If running in a context where an event loop is already running
-            if asyncio.get_event_loop().is_running():
-                asyncio.create_task(self.shutdown_worker()) # schedule it
-            else: # If no loop, run it
-                try:
-                    asyncio.run(self.shutdown_worker())
-                except RuntimeError as e: # Handles "cannot be called when another loop is running"
-                     logger.error(f"RuntimeError during cleanup's shutdown_worker: {e}. May need manual loop management.")
+                logs.append(f"Error with {model_name} on batch: {e}")
+            current_img_offset += len(batch_images)
+        logs.append(f"Finished with {model_name}.")
+
+    # Calculate Final Scores
+    for i in range(num_images):
+        img_scores = [results_data[i][MODEL_REGISTRY[mk]['name']] for mk in selected_model_keys 
+                      if pd.notna(results_data[i].get(MODEL_REGISTRY[mk]['name']))]
+        if img_scores:
+            results_data[i]['Final Score'] = float(np.clip(np.mean(img_scores), 0.0, 10.0))
+
+    df = pd.DataFrame(results_data)
+    # Define column order: Thumbnail, File Name, then model scores, then Final Score
+    ordered_cols = ['Thumbnail', 'File Name'] + \
+                   [MODEL_REGISTRY[k]['name'] for k in MODEL_REGISTRY.keys() if MODEL_REGISTRY[k]['name'] in df.columns] + \
+                   ['Final Score']
+    df = df[[col for col in ordered_cols if col in df.columns]] # Ensure all columns exist
+
+    logs.append("Evaluation complete.")
+    progress_tracker(1.0, desc="Evaluation complete.")
+    return df, logs
+
+def results_df_to_csv_bytes(df: pd.DataFrame, selected_model_display_names: list[str]) -> bytes | None:
+    if df.empty: return None
+    
+    cols_for_csv = ['File Name', 'Final Score'] + \
+                   [name for name in selected_model_display_names if name in df.columns and name not in cols_for_csv]
+    
+    df_csv = df[cols_for_csv].copy()
+    for col in df_csv.select_dtypes(include=['float']).columns: # Format float scores
+        df_csv[col] = df_csv[col].apply(lambda x: f"{x:.4f}" if pd.notnull(x) else "N/A")
+    
+    s_io = io.StringIO()
+    df_csv.to_csv(s_io, index=False)
+    return s_io.getvalue().encode('utf-8')
 
-        # Clean up temporary CSV files
-        for f_path in self_temp_files_to_clean:
-            try:
-                os.remove(f_path)
-                logger.info(f"Removed temp file: {f_path}")
-            except OSError as e:
-                logger.error(f"Error removing temp file {f_path}: {e}")
-        self._temp_files_to_clean.clear()
-        logger.info("Cleanup finished.")
-
-
-#####################################
-#             Interface             #
-#####################################
-
-# Initialize ModelManager once
-model_manager = ModelManager(cache_dir=".model_cache")
-
-def create_interface():
-    # Define model choices based on ModelManager's loaded models
-    # Filter out models that failed to load
-    AVAILABLE_MODEL_KEYS = [k for k in model_manager.model_configs.keys() if k in model_manager.models]
-    AVAILABLE_MODEL_NAMES_MAP = {k: model_manager.model_configs[k]['name'] for k in AVAILABLE_MODEL_KEYS}
+# --- Gradio Interface ---
+def create_gradio_interface():
+    model_name_choices = [config['name'] for config in MODEL_REGISTRY.values()]
     
-    # [(display_name, value_key), ...] for CheckboxGroup
-    MODEL_CHOICES_FOR_CHECKBOX = [(AVAILABLE_MODEL_NAMES_MAP[k], k) for k in AVAILABLE_MODEL_KEYS]
-
-
-    with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky)) as demo:
-        gr.Markdown("""
-        # Comprehensive Image Evaluation Tool (Refactored)
-        Upload images to evaluate them using multiple aesthetic and quality prediction models.
-        Results are displayed in a sortable table with image previews.
-        """)
-
-        # Stores full processing results (list of dicts)
-        # Dict keys: 'file_name', 'final_score', and all model_keys with their scores
-        # This state is the source of truth for regenerating table and CSV
-        results_state = gr.State([]) 
-        # Stores current list of selected model keys (e.g., ['waifu_scorer', 'anime_aesthetic'])
-        selected_models_state = gr.State(AVAILABLE_MODEL_KEYS) 
-        # Stores current log messages as a list
-        log_messages_state = gr.State([])
+    # Define column structure for DataFrame
+    initial_df_cols = ['Thumbnail', 'File Name'] + model_name_choices + ['Final Score']
+    initial_datatypes = ['image', 'str'] + ['number'] * (len(model_name_choices) + 1)
+
+    with gr.Blocks(theme=gr.themes.Glass()) as demo:
+        gr.Markdown("## ✨ Comprehensive Image Evaluation Tool ✨")
+        
+        # For storing results DataFrame between interactions
+        results_state = gr.State(pd.DataFrame(columns=initial_df_cols))
 
         with gr.Row():
-            with gr.Column(scale=1): # Inputs
-                input_images = gr.Files(label="Upload Images", file_count="multiple", type="filepath")
+            with gr.Column(scale=1, min_width=300):
+                gr.Markdown("#### Controls")
+                files_input = gr.Files(label="Upload Images", file_count="multiple", type="filepath")
+                models_checkbox_group = gr.CheckboxGroup(choices=model_name_choices, value=model_name_choices, label="Select Models")
                 
-                if not MODEL_CHOICES_FOR_CHECKBOX:
-                    gr.Markdown("## No models loaded successfully. Please check logs.")
-                    model_checkboxes = None # No models, no checkbox
-                else:
-                    model_checkboxes = gr.CheckboxGroup(
-                        choices=MODEL_CHOICES_FOR_CHECKBOX,
-                        label="Select Models",
-                        value=AVAILABLE_MODEL_KEYS, # Default to all available selected
-                        info="Choose models for evaluation. Final score is an average of selected model scores."
-                    )
-
-                auto_batch_checkbox = gr.Checkbox(label="Automatic Batch Size Detection", value=True)
-                batch_size_input = gr.Number(label="Manual Batch Size", value=8, minimum=1, precision=0, interactive=False) # Interactive based on auto_batch_checkbox
+                with gr.Accordion("Batch Settings", open=False):
+                    auto_batch_toggle = gr.Checkbox(label="Auto-detect Batch Size", value=True)
+                    manual_batch_input = gr.Number(label="Manual Batch Size", value=4, minimum=1, step=1, interactive=False) # Interactive based on toggle
                 
-                process_btn = gr.Button("Evaluate Images", variant="primary", interactive=bool(MODEL_CHOICES_FOR_CHECKBOX))
-                clear_btn = gr.Button("Clear Results")
-                download_csv_btn = gr.Button("Download Results as CSV", variant="secondary")
-
-            with gr.Column(scale=3): # Outputs
-                progress_tracker = gr.Progress(label="Processing Progress")
-                log_output = gr.Textbox(label="Logs", lines=10, max_lines=20, interactive=False, autoscroll=True)
+                evaluate_button = gr.Button("🚀 Evaluate Images", variant="primary")
+                with gr.Row():
+                    clear_button = gr.Button("🧹 Clear")
+                    download_button = gr.Button("💾 Download CSV")
+                
+                # Hidden component for file download functionality
+                csv_file_output = gr.File(label="Download CSV File", visible=False) 
+            
+            with gr.Column(scale=3, min_width=600):
+                gr.Markdown("#### Results")
+                # Using gr.Slider for progress display
+                progress_slider = gr.Slider(label="Progress", minimum=0, maximum=1, value=0, interactive=False)
                 
-                # Initial headers for DataFrame; will be updated dynamically
-                initial_df_headers = ['Image', 'File Name'] + [AVAILABLE_MODEL_NAMES_MAP[k] for k in AVAILABLE_MODEL_KEYS] + ['Final Score']
                 results_dataframe = gr.DataFrame(
-                    headers=initial_df_headers,
-                    datatype=['pil'] + ['str'] * (len(initial_df_headers) -1) , # Image + strings for scores
-                    label="Evaluation Results",
-                    interactive=True, # Enables sorting by clicking headers
-                    row_count=(10, "dynamic"),
-                    col_count=(len(initial_df_headers), "fixed"),
-                    wrap=True,
+                    label="Evaluation Scores",
+                    headers=initial_df_cols,
+                    datatype=initial_datatypes,
+                    interactive=True, # Enables native sorting by clicking headers
+                    height=500,
+                    wrap=True
                 )
-                # Hidden file component for download trigger
-                download_file_provider = gr.File(label="Download Link", visible=False)
-
-        # --- Callback Functions ---
-        def update_batch_size_interactive(auto_detect_enabled: bool):
-            return gr.Number.update(interactive=not auto_detect_enabled)
-
-        async def handle_process_images_ui(
-            files_list: list[gr. rýchle.TempFile] | None, # Gradio File objects
-            auto_batch_flag: bool, 
-            manual_batch_val: int, 
-            selected_model_keys_from_ui: list[str],
-            # Gradio will pass the gr.Progress instance automatically by type hinting
-            # Ensure the name 'progress_tracker_instance' matches an output component if you want to update it by dict key
-            # Otherwise, use the positional argument `progress`
-            progress_instance: gr.Progress 
-        ):
-            if not files_list:
-                yield {
-                    log_output: "No files uploaded. Please select images first.",
-                    progress_tracker: gr.Progress(0.0, "Idle. No files."),
-                    results_dataframe: gr.DataFrame.update(value=None), # Clear table
-                    results_state: [],
-                    selected_models_state: selected_model_keys_from_ui,
-                    log_messages_state: ["No files uploaded. Please select images first."]
-                }
-                return
+                logs_textbox = gr.Textbox(label="Process Logs", lines=5, max_lines=10, interactive=False)
 
-            # Update selected_models_state right away
-            yield { selected_models_state: selected_model_keys_from_ui, log_messages_state: [] } # Clear logs state
+        # --- Callbacks ---
+        def map_display_names_to_keys(display_names: list[str]) -> list[str]:
+            return [key for key, cfg in MODEL_REGISTRY.items() if cfg['name'] in display_names]
 
-            # Convert Gradio TempFile objects to string paths
-            actual_file_paths = [f.name for f in files_list]
+        async def run_evaluation(uploaded_files, selected_model_names, auto_batch, manual_batch, 
+                                 current_results_df, progress=gr.Progress(track_tqdm=True)):
+            if not uploaded_files:
+                return {
+                    results_state: current_results_df, logs_textbox: "No files uploaded. Please upload images first.",
+                    progress_slider: gr.update(value=0, label="Progress")
+                }
             
-            current_log_list = [] # Local log accumulator for this run
-
-            # Call the ModelManager's generator
-            # The progress_instance is implicitly passed by Gradio to this function
-            # The ModelManager generator will then use it via its own parameter `progress_tracker_instance`
+            yield {logs_textbox: "Loading images...", progress_slider: gr.update(value=0.01, label="Loading images...")}
             
-            # Need to get an async generator from model_manager
-            evaluation_generator = await model_manager.submit_evaluation_request(
-                actual_file_paths, auto_batch_flag, manual_batch_val, selected_model_keys_from_ui
-            )
+            pil_images, file_names = [], []
+            for f_obj in uploaded_files:
+                try:
+                    pil_images.append(Image.open(f_obj.name).convert("RGB")) # f_obj.name is path for type="filepath"
+                    file_names.append(os.path.basename(f_obj.name))
+                except Exception as e:
+                    print(f"Error loading image {f_obj.name}: {e}") # Log to console
             
-            dataframe_update_value = None
-            final_results_for_app_state = []
-
-            async for event in evaluation_generator:
-                outputs_to_yield = {}
-                if event["type"] == "log_update":
-                    current_log_list = event["messages"]
-                    outputs_to_yield[log_output] = "\n".join(current_log_list)
-                elif event["type"] == "progress":
-                    # Update progress bar directly using the passed instance
-                    progress_instance(event["value"], desc=event.get("desc")) 
-                elif event["type"] == "batch_size_update":
-                    outputs_to_yield[batch_size_input] = gr.Number.update(value=event["value"])
-                elif event["type"] == "partial_results_df_rows":
-                    # data is list of lists for DataFrame rows
-                    # selected_model_keys used to generate current headers
-                    dynamic_headers = ['Image', 'File Name'] + \
-                                      [AVAILABLE_MODEL_NAMES_MAP[k] for k in event["selected_model_keys"] if k in AVAILABLE_MODEL_NAMES_MAP] + \
-                                      ['Final Score']
-                    dataframe_update_value = pd.DataFrame(event["data"], columns=dynamic_headers) if event["data"] else None
-                    outputs_to_yield[results_dataframe] = gr.DataFrame.update(value=dataframe_update_value, headers=dynamic_headers)
-
-                elif event["type"] == "final_results_state":
-                    final_results_for_app_state = event["data"]
-                
-                if outputs_to_yield: # Only yield if there's something to update
-                     yield outputs_to_yield
+            if not pil_images:
+                return {logs_textbox: "No valid images could be loaded.", progress_slider: gr.update(value=0, label="Error")}
+
+            selected_keys = map_display_names_to_keys(selected_model_names)
             
-            # Final updates after generator is exhausted
-            yield {
-                results_state: final_results_for_app_state,
-                log_messages_state: current_log_list, # Save final logs
-                # DataFrame should be up-to-date from the last partial_results_df_rows
-            }
+            batch_size_to_use = manual_batch
+            if auto_batch:
+                yield {logs_textbox: "Auto-tuning batch size...", progress_slider: gr.update(value=0.1, label="Auto-tuning...")}
+                batch_size_to_use = auto_tune_batch_size(pil_images, selected_keys, verbose=True)
+                yield {manual_batch_input: gr.update(value=batch_size_to_use)} # Update UI with detected size
+            
+            yield {logs_textbox: f"Starting evaluation with batch size {batch_size_to_use}...", 
+                   progress_slider: gr.update(value=0.15, label=f"Evaluating (Batch: {batch_size_to_use})...")}
 
+            df_new_results, log_messages = await evaluate_images_core(
+                pil_images, file_names, selected_keys, batch_size_to_use, progress
+            )
+            
+            # Sort by 'Final Score' descending by default before display
+            if not df_new_results.empty and 'Final Score' in df_new_results.columns:
+                df_new_results = df_new_results.sort_values(by='Final Score', ascending=False, na_position='last')
 
-        def handle_clear_results_ui():
-            # Clear files, logs, table, progress, and internal states
             return {
-                input_images: None,
-                log_output: "Results cleared.",
-                results_dataframe: gr.DataFrame.update(value=None, headers=initial_df_headers), # Reset with initial headers
-                progress_tracker: gr.Progress(0.0, "Idle"),
-                results_state: [],
-                # selected_models_state: AVAILABLE_MODEL_KEYS, # Optionally reset model selection
-                batch_size_input: gr.Number.update(value=8), # Reset batch size
-                log_messages_state: ["Results cleared."]
+                results_state: df_new_results, results_dataframe: df_new_results,
+                logs_textbox: "\n".join(log_messages),
+                progress_slider: gr.update(value=1.0, label="Evaluation Complete")
             }
-        
-        # Function to re-render DataFrame and update states when model selection changes
-        def handle_model_selection_or_state_change_ui(
-            current_selected_keys: list[str], 
-            current_full_results: list[dict]
-        ):
-            if not current_full_results: # No data to process
-                dynamic_headers = ['Image', 'File Name'] + \
-                                  [AVAILABLE_MODEL_NAMES_MAP[k] for k in current_selected_keys if k in AVAILABLE_MODEL_NAMES_MAP] + \
-                                  ['Final Score']
-                return {
-                    results_dataframe: gr.DataFrame.update(value=None, headers=dynamic_headers),
-                    selected_models_state: current_selected_keys,
-                    results_state: current_full_results # pass through if empty
-                }
-
-            new_df_rows = []
-            updated_full_results = []
 
-            for res_item_dict in current_full_results:
-                # Recalculate final score based on new selection
-                scores_to_avg = []
-                for mk in current_selected_keys:
-                    if mk in res_item_dict and isinstance(res_item_dict[mk], (float, int)):
-                        scores_to_avg.append(res_item_dict[mk])
-                
-                new_final_score = None
-                if scores_to_avg:
-                    new_final_score_val = float(np.mean(scores_to_avg))
-                    new_final_score = float(np.clip(new_final_score_val, 0.0, 10.0))
-                
-                # Update the item in results_state
-                res_item_dict['final_score'] = new_final_score
-                updated_full_results.append(res_item_dict.copy()) # Store updated item
-
-                # Prepare row for DataFrame
-                # Find the corresponding image (this assumes images are not stored in results_state, which they aren't)
-                # For simplicity, we'll need to re-generate thumbnails if we want them in this update path.
-                # A robust way: results_state stores paths or minimal data to re-fetch/re-create thumbnails.
-                # Current implementation of `run_evaluation_generator` directly yields DF rows with PIL images.
-                # If `handle_model_selection_change_ui` is to re-generate the DF from `results_state`,
-                # `results_state` items would need to include enough info for `Image.open` and `thumbnail`.
-                # This is a complex part if we want perfect dynamic DF regeneration with images.
-                # For now, let's assume `results_state` stores `PIL.Image` thumbnails if this path is critical.
-                # The `run_evaluation_generator` stores dicts without PIL image objects in `all_results_for_state`.
-                # This means `handle_model_selection_change_ui` cannot easily reconstruct the 'Image' column.
-                #
-                # SIMPLIFICATION: When model selection changes, we only update scores in the existing DataFrame
-                # if possible, or we re-calculate and re-populate. The current code path re-creates rows.
-                # To do this properly, `results_state` items should perhaps include original image path or cached thumbnail.
-                #
-                # Let's make results_state store {'file_path': ..., 'thumbnail_pil': ..., scores...}
-                # This needs `run_evaluation_generator` to save file_path and thumbnail_pil to `all_results_for_state`.
-                # Assume `results_state` items now contain 'thumbnail_pil' and other scores.
-                
-                # If 'thumbnail_pil' is not in res_item_dict (because it wasn't saved that way), this will fail.
-                # This path requires results_state to contain PIL image data for the 'Image' column.
-                # The current 'run_evaluation_generator' does not save PIL images into `all_results_for_state`.
-                # It only creates them for immediate DataFrame update.
-                # This function needs to be re-thought if full DF reconstruction with images is needed here.
-
-                # Let's assume results_state IS NOT used to rebuild the image column.
-                # The change handler for model_checkboxes will mostly affect the *calculation* of final_score
-                # and *visibility* of columns if we were dynamically adding/removing them.
-                # Gradio's DataFrame doesn't easily hide/show columns; we change headers and data.
-                
-                # Rebuild row for DF:
-                df_row = [res_item_dict.get('thumbnail_pil_placeholder', "N/A"), res_item_dict['file_name']]
-                for mk_cfg in AVAILABLE_MODEL_KEYS: # All possible models to maintain column order
-                    if mk_cfg in current_selected_keys: # If this model is currently selected for display
-                        score = res_item_dict.get(mk_cfg)
-                        df_row.append(f"{score:.4f}" if isinstance(score, (float, int)) else "N/A")
-                    # If not selected, this column won't even be in dynamic_headers.
-                df_row.append(f"{new_final_score:.4f}" if new_final_score is not None else "N/A")
-                new_df_rows.append(df_row)
-            
-            dynamic_headers = ['Image', 'File Name'] + \
-                              [AVAILABLE_MODEL_NAMES_MAP[k] for k in current_selected_keys if k in AVAILABLE_MODEL_NAMES_MAP] + \
-                              ['Final Score']
-            
-            import pandas as pd
-            df_value = pd.DataFrame(new_df_rows, columns=dynamic_headers) if new_df_rows else None
-            
+        def clear_all_outputs():
+            empty_df = pd.DataFrame(columns=initial_df_cols)
             return {
-                results_dataframe: gr.DataFrame.update(value=df_value, headers=dynamic_headers),
-                selected_models_state: current_selected_keys, # Persist the new selection
-                results_state: updated_full_results # Persist updated scores
+                results_state: empty_df, results_dataframe: empty_df,
+                files_input: None, logs_textbox: "Outputs cleared.",
+                progress_slider: gr.update(value=0, label="Progress")
             }
 
-
-        def handle_download_csv_ui(current_full_results: list[dict], current_selected_keys: list[str]):
-            if not current_full_results:
-                # Optionally, send a message to log_output if desired using yield
-                return gr.File.update(value=None, visible=False) 
-
-            # Use StringIO to build CSV in memory
-            csv_output = StringIO()
-            # Define fieldnames: Filename, selected model scores, Final Score
-            fieldnames = ['File Name'] + \
-                         [AVAILABLE_MODEL_NAMES_MAP[k] for k in current_selected_keys if k in AVAILABLE_MODEL_NAMES_MAP] + \
-                         ['Final Score']
+        def download_csv_file(current_df, selected_names):
+            if current_df.empty:
+                gr.Warning("No results available to download.")
+                return None
             
-            writer = csv.DictWriter(csv_output, fieldnames=fieldnames, extrasaction='ignore')
-            writer.writeheader()
-
-            for res_item in current_full_results:
-                row_to_write = {'File Name': res_item['file_name']}
-                final_score_val = res_item.get('final_score') # This should be up-to-date from results_state
-                row_to_write['Final Score'] = f"{final_score_val:.4f}" if final_score_val is not None else "N/A"
-                
-                for key in current_selected_keys:
-                    if key in AVAILABLE_MODEL_NAMES_MAP: # ensure it's a valid model key
-                        model_display_name = AVAILABLE_MODEL_NAMES_MAP[key]
-                        score_val = res_item.get(key)
-                        row_to_write[model_display_name] = f"{score_val:.4f}" if isinstance(score_val, (float, int)) else "N/A"
-                writer.writerow(row_to_write)
+            csv_data = results_df_to_csv_bytes(current_df, selected_names)
+            if csv_data:
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".csv", mode='wb') as tmp_f:
+                    tmp_f.write(csv_data)
+                gr.Info("CSV file prepared for download.")
+                return tmp_f.name
+            gr.Error("Failed to generate CSV.")
+            return None
+        
+        def update_final_scores_on_model_select(selected_model_names, current_df):
+            if current_df.empty: return current_df
             
-            csv_content = csv_output.getvalue()
-            csv_output.close()
-
-            # Save to a temporary file that Gradio can serve
-            with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".csv", encoding='utf-8') as tmp_file:
-                tmp_file.write(csv_content)
-                temp_file_path = tmp_file.name
+            df_updated = current_df.copy()
+            selected_keys = map_display_names_to_keys(selected_model_names)
+
+            for i, row in df_updated.iterrows():
+                img_scores = [row[MODEL_REGISTRY[mk]['name']] for mk in selected_keys 
+                              if pd.notna(row.get(MODEL_REGISTRY[mk]['name']))]
+                if img_scores:
+                    df_updated.loc[i, 'Final Score'] = float(np.clip(np.mean(img_scores), 0.0, 10.0))
+                else:
+                    df_updated.loc[i, 'Final Score'] = np.nan
             
-            model_manager.add_temp_file_for_cleanup(temp_file_path) # Register for cleanup
-
-            return gr.File.update(value=temp_file_path, visible=True, label="results.csv")
-
-
-        # --- Wire up components ---
-        auto_batch_checkbox.change(
-            fn=update_batch_size_interactive,
-            inputs=[auto_batch_checkbox],
-            outputs=[batch_size_input]
-        )
+            if 'Final Score' in df_updated.columns: # Re-sort
+                df_updated = df_updated.sort_values(by='Final Score', ascending=False, na_position='last')
 
-        # Check if model_checkboxes exists (i.e., models loaded)
-        if model_checkboxes:
-            process_btn.click(
-                fn=handle_process_images_ui,
-                inputs=[input_images, auto_batch_checkbox, batch_size_input, model_checkboxes],
-                outputs=[
-                    log_output, progress_tracker, results_dataframe, batch_size_input, 
-                    results_state, selected_models_state, log_messages_state # Ensure all yielded components are listed
-                ]
-            )
-            # When model selection changes, update the displayed table and internal states
-            model_checkboxes.change(
-                fn=handle_model_selection_or_state_change_ui,
-                inputs=[model_checkboxes, results_state], # Takes current selection and full results data
-                outputs=[results_dataframe, selected_models_state, results_state]
-            )
+            return {results_state: df_updated, results_dataframe: df_updated}
 
-        clear_btn.click(
-            fn=handle_clear_results_ui,
-            outputs=[
-                input_images, log_output, results_dataframe, progress_tracker, 
-                results_state, batch_size_input, log_messages_state # model_checkboxes could be reset too if needed
-            ]
-        )
+        auto_batch_toggle.change(lambda x: gr.update(interactive=not x), inputs=auto_batch_toggle, outputs=manual_batch_input)
         
-        download_csv_btn.click(
-            fn=handle_download_csv_ui,
-            inputs=[results_state, selected_models_state], # Use current results and selected models for CSV
-            outputs=[download_file_provider]
+        evaluate_button.click(
+            fn=run_evaluation,
+            inputs=[files_input, models_checkbox_group, auto_batch_toggle, manual_batch_input, results_state],
+            outputs=[results_state, results_dataframe, logs_textbox, manual_batch_input, progress_slider]
         )
-
-        # Initial setup on demo load
-        async def initial_load_setup():
-            await model_manager.start_worker_if_not_running() # Start async worker
-            # Set initial state for selected_models_state based on default checkbox values
-            # This is a bit of a workaround if direct binding isn't available for initial state from component value
-            return {selected_models_state: AVAILABLE_MODEL_KEYS, log_messages_state: ["Application loaded. Ready."]}
-
-        demo.load(
-            fn=initial_load_setup,
-            outputs=[selected_models_state, log_messages_state]
+        clear_button.click(fn=clear_all_outputs, outputs=[results_state, results_dataframe, files_input, logs_textbox, progress_slider])
+        download_button.click(fn=download_csv_file, inputs=[results_state, models_checkbox_group], outputs=csv_file_output)
+        models_checkbox_group.change(
+            fn=update_final_scores_on_model_select,
+            inputs=[models_checkbox_group, results_state],
+            outputs=[results_state, results_dataframe]
         )
-        # Register cleanup function
-        demo.unload(model_manager.cleanup)
-
-
-        gr.Markdown("""
-        ### Notes
-        - **Model Selection**: Dynamically choose models for evaluation. The 'Final Score' and displayed columns update accordingly.
-        - **Native Table**: Results are shown in a native Gradio DataFrame, allowing sorting by clicking column headers.
-        - **Batching**: Automatic batch size detection is enabled by default. You can switch to manual batch sizing.
-        - **CSV Export**: Download the current results (respecting selected models for columns) as a CSV file.
-        - **Asynchronous Processing**: Image evaluation runs in the background, providing live updates for logs and progress.
-        """)
-    return demo
 
+        # Initial load state for the DataFrame UI component
+        demo.load(lambda: pd.DataFrame(columns=initial_df_cols), outputs=[results_dataframe])
+    return demo
 
 if __name__ == "__main__":
-    # For proper MLP Safetensor loading with adjusted keys, ensure 'safetensors' is installed.
-    # For WaifuScorer, ensure 'openai-clip' is installed.
-    # For ONNX models, 'onnxruntime' or 'onnxruntime-gpu'.
-    # For general Hugging Face model loading, 'transformers'.
-    # OpenCV ('cv2') for image manipulation: 'opencv-python'.
-    # And of course 'torch', 'numpy', 'Pillow', 'gradio'.
-
-    # Create a dummy aesthetic_predictor_v2_5.py if it doesn't exist for the stub to work
-    # (or ensure the real one is present)
-    if not Path("aesthetic_predictor_v2_5.py").exists():
-        stub_content = """
-# Placeholder for aesthetic_predictor_v2_5.py
-# This file needs to contain the actual 'convert_v2_5_from_siglip' function.
-# The main script uses a basic stub if this file is missing or fails to import.
-# print("aesthetic_predictor_v2_5.py placeholder executed")
-def convert_v2_5_from_siglip(*args, **kwargs):
-    raise NotImplementedError("This is a placeholder. Implement convert_v2_5_from_siglip here or ensure the main script's stub is used.")
-"""
-        # Only write if you are sure, or better, let user handle this dependency.
-        # For this exercise, we assume the main script's internal stub is sufficient if the file is missing.
-        pass
-
-
-    # It's important that the ModelManager is initialized before create_interface() is called,
-    # as create_interface() relies on model_manager.model_configs.
-    # This is already handled by placing `model_manager = ModelManager()` globally.
-    
-    app_interface = create_interface()
-    app_interface.queue().launch(debug=True, share=False) # Enable queue for async operations
-
-    # Ensure cleanup is called on exit if demo.unload isn't fully effective in all environments
-    import atexit
-    atexit.register(model_manager.cleanup)
+    initialize_models(verbose_loading=True) # Load models once at startup
+    gradio_app = create_gradio_interface()
+    gradio_app.queue().launch(debug=False) # Enable queue for async ops, debug=True for more logs