Update app.py

app.py

@@ -1,710 +1,824 @@
+"""
+Modern Image Evaluation Tool with Aesthetic and Quality Prediction Models
+
+This refactored version features:
+- Modern async/await patterns with proper error handling
+- Type hints throughout for better code maintainability
+- Dependency injection and factory patterns
+- Proper resource management with context managers
+- Configuration-driven model loading
+- Improved batch processing with memory optimization
+- Clean separation of concerns with proper abstraction layers
+"""
+
+import asyncio
+import base64
+import csv
+import logging
 import os
-import shutil
 import tempfile
-import base64
-import asyncio
-from io import BytesIO
+import shutil
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from enum import Enum
+from io import BytesIO, StringIO
+from pathlib import Path
+from typing import Dict, List, Optional, Protocol, Tuple, Union, Any
+from abc import ABC, abstractmethod
 
 import cv2
+import gradio as gr
 import numpy as np
+import onnxruntime as ort
 import torch
-import onnxruntime as rt
+import torch.nn as nn
 from PIL import Image
-import gradio as gr
 from transformers import pipeline
 from huggingface_hub import hf_hub_download
 
-# 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):
-    """A simple multi-layer perceptron for image feature regression."""
-    def __init__(self, input_size: int, batch_norm: bool = True):
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Configuration and Data Models
+# =============================================================================
+
+class ModelType(Enum):
+    """Enumeration of available model types."""
+    AESTHETIC_SHADOW = "aesthetic_shadow"
+    WAIFU_SCORER = "waifu_scorer"
+    AESTHETIC_PREDICTOR_V2_5 = "aesthetic_predictor_v2_5"
+    ANIME_AESTHETIC = "anime_aesthetic"
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for individual models."""
+    name: str
+    display_name: str
+    enabled: bool = True
+    batch_supported: bool = True
+    model_path: Optional[str] = None
+    cache_dir: Optional[str] = None
+
+
+@dataclass
+class ProcessingConfig:
+    """Configuration for processing parameters."""
+    auto_batch: bool = False
+    manual_batch_size: int = 1
+    max_batch_size: int = 64
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    score_range: Tuple[float, float] = (0.0, 10.0)
+
+
+@dataclass
+class EvaluationResult:
+    """Data class for individual evaluation results."""
+    file_name: str
+    file_path: str
+    thumbnail_b64: str
+    model_scores: Dict[str, Optional[float]] = field(default_factory=dict)
+    final_score: Optional[float] = None
+    processing_time: float = 0.0
+    error: Optional[str] = None
+
+
+@dataclass
+class BatchResult:
+    """Data class for batch processing results."""
+    results: List[EvaluationResult]
+    logs: List[str]
+    processing_time: float
+    batch_size_used: int
+    success_count: int
+    error_count: int
+
+
+# =============================================================================
+# Model Interfaces and Implementations
+# =============================================================================
+
+class BaseModel(Protocol):
+    """Protocol defining the interface for all evaluation models."""
+    
+    async def predict(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Predict scores for a batch of images."""
+        ...
+    
+    def is_available(self) -> bool:
+        """Check if the model is available and ready for inference."""
+        ...
+    
+    def cleanup(self) -> None:
+        """Clean up model resources."""
+        ...
+
+
+class ModernMLP(nn.Module):
+    """Modern implementation of MLP with improved architecture."""
+    
+    def __init__(
+        self, 
+        input_size: int, 
+        hidden_dims: List[int] = None,
+        dropout_rates: List[float] = None,
+        use_batch_norm: bool = True,
+        activation: nn.Module = nn.ReLU
+    ):
         super().__init__()
-        self.input_size = input_size
-        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(),
-            torch.nn.BatchNorm1d(512) if batch_norm else torch.nn.Identity(),
-            torch.nn.Dropout(0.3),
-            torch.nn.Linear(512, 256),
-            torch.nn.ReLU(),
-            torch.nn.BatchNorm1d(256) if batch_norm else torch.nn.Identity(),
-            torch.nn.Dropout(0.2),
-            torch.nn.Linear(256, 128),
-            torch.nn.ReLU(),
-            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)
-        )
-
+        
+        if hidden_dims is None:
+            hidden_dims = [2048, 512, 256, 128, 32]
+        if dropout_rates is None:
+            dropout_rates = [0.3, 0.3, 0.2, 0.1, 0.0]
+            
+        layers = []
+        prev_dim = input_size
+        
+        for i, (hidden_dim, dropout_rate) in enumerate(zip(hidden_dims, dropout_rates)):
+            layers.append(nn.Linear(prev_dim, hidden_dim))
+            layers.append(activation())
+            
+            if use_batch_norm and i < len(hidden_dims) - 1:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            
+            if dropout_rate > 0:
+                layers.append(nn.Dropout(dropout_rate))
+                
+            prev_dim = hidden_dim
+        
+        # Final output layer
+        layers.append(nn.Linear(prev_dim, 1))
+        self.network = nn.Sequential(*layers)
+    
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.layers(x)
+        return self.network(x)
 
 
-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
+class WaifuScorerModel:
+    """Modernized WaifuScorer implementation with better error handling."""
+    
+    def __init__(self, config: ModelConfig, device: str):
+        self.config = config
         self.device = device
         self.dtype = torch.float32
-        self.available = False
-
+        self._available = False
+        self._model = None
+        self._clip_model = None
+        self._preprocess = None
+        
+        self._initialize_model()
+    
+    def _initialize_model(self) -> None:
+        """Initialize the model with proper error handling."""
         try:
-            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 provided. Using default: {model_path}")
-
-            # Download model if not found locally
-            if not os.path.isfile(model_path):
-                username, repo_id, model_name = model_path.split("/")[-3:]
-                model_path = hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=cache_dir)
-
-            if self.verbose:
-                print(f"Loading WaifuScorer model from: {model_path}")
-
-            # Initialize MLP model
-            self.mlp = MLP(input_size=768)
-            # Load state dict
+            import clip
+            
+            # Download model if needed
+            model_path = self._get_model_path()
+            
+            # Initialize MLP
+            self._model = ModernMLP(input_size=768)
+            
+            # Load weights
             if model_path.endswith(".safetensors"):
                 from safetensors.torch import load_file
                 state_dict = load_file(model_path)
             else:
-                state_dict = torch.load(model_path, map_location=device)
-            self.mlp.load_state_dict(state_dict)
-            self.mlp.to(device)
-            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
+                state_dict = torch.load(model_path, map_location=self.device)
+            
+            self._model.load_state_dict(state_dict)
+            self._model.to(self.device)
+            self._model.eval()
+            
+            # Load CLIP model
+            self._clip_model, self._preprocess = clip.load("ViT-L/14", device=self.device)
+            self._available = True
+            
+            logger.info(f"WaifuScorer model loaded successfully on {self.device}")
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize WaifuScorer: {e}")
+            self._available = False
+    
+    def _get_model_path(self) -> str:
+        """Get or download the model path."""
+        if self.config.model_path and os.path.isfile(self.config.model_path):
+            return self.config.model_path
+        
+        # Default download path
+        model_path = "Eugeoter/waifu-scorer-v3/model.pth"
+        username, repo_id, model_name = model_path.split("/")[-3:]
+        return hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=self.config.cache_dir)
+    
+    async def predict(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Predict scores for a batch of images."""
+        if not self._available:
+            return [None] * len(images)
+        
+        try:
+            # Handle single image case for CLIP compatibility
+            batch_images = images * 2 if len(images) == 1 else images
+            
+            # Preprocess images
+            image_tensors = [self._preprocess(img).unsqueeze(0) for img in batch_images]
+            image_batch = torch.cat(image_tensors).to(self.device)
+            
+            # Extract features and predict
+            with torch.no_grad():
+                image_features = self._clip_model.encode_image(image_batch)
+                # Normalize features
+                norm = image_features.norm(2, dim=-1, keepdim=True)
+                norm[norm == 0] = 1
+                normalized_features = (image_features / norm).to(device=self.device, dtype=self.dtype)
+                
+                predictions = self._model(normalized_features)
+                scores = predictions.clamp(0, 10).cpu().numpy().reshape(-1).tolist()
+                
+            return scores[:len(images)]
+            
+        except Exception as e:
+            logger.error(f"Error in WaifuScorer prediction: {e}")
+            return [None] * len(images)
+    
+    def is_available(self) -> bool:
+        return self._available
+    
+    def cleanup(self) -> None:
+        """Clean up model resources."""
+        if self._model is not None:
+            del self._model
+        if self._clip_model is not None:
+            del self._clip_model
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+
+class AestheticShadowModel:
+    """Wrapper for Aesthetic Shadow model using transformers pipeline."""
+    
+    def __init__(self, config: ModelConfig, device: str):
+        self.config = config
+        self.device = device
+        self._available = False
+        self._model = None
+        
+        self._initialize_model()
+    
+    def _initialize_model(self) -> None:
+        """Initialize the model pipeline."""
+        try:
+            self._model = pipeline(
+                "image-classification", 
+                model="NeoChen1024/aesthetic-shadow-v2-backup", 
+                device=self.device
+            )
+            self._available = True
+            logger.info("Aesthetic Shadow model loaded successfully")
+            
         except Exception as e:
-            print(f"Unable to initialize WaifuScorer: {e}")
-
-    @torch.no_grad()
-    def __call__(self, images):
-        if not self.available:
-            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
-
-        image_tensors = [self.preprocess(img).unsqueeze(0) for img in images]
-        image_batch = torch.cat(image_tensors).to(self.device)
-        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():
-    """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(
+            logger.error(f"Failed to initialize Aesthetic Shadow: {e}")
+            self._available = False
+    
+    async def predict(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Predict scores for a batch of images."""
+        if not self._available:
+            return [None] * len(images)
+        
+        try:
+            results = self._model(images)
+            scores = []
+            
+            for result in results:
+                try:
+                    hq_score = next(p for p in result if p['label'] == 'hq')['score']
+                    score = float(np.clip(hq_score * 10.0, 0.0, 10.0))
+                    scores.append(score)
+                except (StopIteration, KeyError, TypeError):
+                    scores.append(None)
+            
+            return scores
+            
+        except Exception as e:
+            logger.error(f"Error in Aesthetic Shadow prediction: {e}")
+            return [None] * len(images)
+    
+    def is_available(self) -> bool:
+        return self._available
+    
+    def cleanup(self) -> None:
+        if self._model is not None:
+            del self._model
+
+
+class AestheticPredictorV25Model:
+    """Wrapper for Aesthetic Predictor V2.5 model."""
+    
+    def __init__(self, config: ModelConfig, device: str):
+        self.config = config
+        self.device = device
+        self._available = False
+        self._model = None
+        self._preprocessor = None
+        
+        self._initialize_model()
+    
+    def _initialize_model(self) -> None:
+        """Initialize the model."""
+        try:
+            from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip
+            
+            self._model, self._preprocessor = convert_v2_5_from_siglip(
                 low_cpu_mem_usage=True,
                 trust_remote_code=True,
             )
+            
             if torch.cuda.is_available():
-                self.model = self.model.to(torch.bfloat16).cuda()
-
-        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 AestheticPredictorV2_5_Impl()
-
-
-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")
-    return rt.InferenceSession(model_path, providers=['CPUExecutionProvider'])
-
-
-def predict_anime_aesthetic(img, model):
-    """Predict Anime Aesthetic score for a single image."""
-    img_np = np.array(img).astype(np.float32) / 255.0
-    s = 768
-    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        #
-#####################################
+                self._model = self._model.to(torch.bfloat16).cuda()
+            
+            self._available = True
+            logger.info("Aesthetic Predictor V2.5 loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize Aesthetic Predictor V2.5: {e}")
+            self._available = False
+    
+    async def predict(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Predict scores for a batch of images."""
+        if not self._available:
+            return [None] * len(images)
+        
+        try:
+            rgb_images = [img.convert("RGB") for img in images]
+            pixel_values = self._preprocessor(images=rgb_images, 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 [float(np.round(np.clip(s, 0.0, 10.0), 4)) for s in scores]
+            
+        except Exception as e:
+            logger.error(f"Error in Aesthetic Predictor V2.5 prediction: {e}")
+            return [None] * len(images)
+    
+    def is_available(self) -> bool:
+        return self._available
+    
+    def cleanup(self) -> None:
+        if self._model is not None:
+            del self._model
+
+
+class AnimeAestheticModel:
+    """ONNX-based Anime Aesthetic model."""
+    
+    def __init__(self, config: ModelConfig, device: str):
+        self.config = config
+        self.device = device
+        self._available = False
+        self._session = None
+        
+        self._initialize_model()
+    
+    def _initialize_model(self) -> None:
+        """Initialize the ONNX model."""
+        try:
+            model_path = hf_hub_download(repo_id="skytnt/anime-aesthetic", filename="model.onnx")
+            self._session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
+            self._available = True
+            logger.info("Anime Aesthetic model loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize Anime Aesthetic: {e}")
+            self._available = False
+    
+    async def predict(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Predict scores for images (single image processing for ONNX)."""
+        if not self._available:
+            return [None] * len(images)
+        
+        scores = []
+        for img in images:
+            try:
+                score = self._predict_single(img)
+                scores.append(float(np.clip(score * 10.0, 0.0, 10.0)))
+            except Exception as e:
+                logger.error(f"Error predicting anime aesthetic for image: {e}")
+                scores.append(None)
+        
+        return scores
+    
+    def _predict_single(self, img: Image.Image) -> float:
+        """Predict score for a single image."""
+        img_np = np.array(img).astype(np.float32) / 255.0
+        s = 768
+        h, w = img_np.shape[:2]
+        
+        # Resize while maintaining aspect ratio
+        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 crop/pad to square
+        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
+        input_tensor = np.transpose(canvas, (2, 0, 1))[np.newaxis, :]
+        return self._session.run(None, {"img": input_tensor})[0].item()
+    
+    def is_available(self) -> bool:
+        return self._available
+    
+    def cleanup(self) -> None:
+        if self._session is not None:
+            del self._session
+
+
+# =============================================================================
+# Model Factory and Manager
+# =============================================================================
+
+class ModelFactory:
+    """Factory for creating model instances."""
+    
+    _MODEL_CLASSES = {
+        ModelType.AESTHETIC_SHADOW: AestheticShadowModel,
+        ModelType.WAIFU_SCORER: WaifuScorerModel,
+        ModelType.AESTHETIC_PREDICTOR_V2_5: AestheticPredictorV25Model,
+        ModelType.ANIME_AESTHETIC: AnimeAestheticModel,
+    }
+    
+    @classmethod
+    def create_model(cls, model_type: ModelType, config: ModelConfig, device: str) -> BaseModel:
+        """Create a model instance based on type."""
+        model_class = cls._MODEL_CLASSES.get(model_type)
+        if not model_class:
+            raise ValueError(f"Unknown model type: {model_type}")
+        
+        return model_class(config, device)
+
 
 class ModelManager:
-    """Manages model loading and processing requests using a queue."""
-    def __init__(self):
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        print(f"Using device: {self.device}")
-        print("Loading models... This may take some time.")
-
-        # Load models once during initialization
-        print("Loading Aesthetic Shadow model...")
-        self.aesthetic_shadow_model = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=self.device)
-        print("Loading Waifu Scorer model...")
-        self.waifu_scorer_model = WaifuScorer(device=self.device, verbose=True)
-        print("Loading Aesthetic Predictor V2.5...")
-        self.aesthetic_predictor_model = load_aesthetic_predictor_v2_5()
-        print("Loading Anime Aesthetic model...")
-        self.anime_aesthetic_model = load_anime_aesthetic_model()
-        print("All models loaded successfully!")
-
-        self.available_models = {
-            "aesthetic_shadow": {"name": "Aesthetic Shadow", "process": self._process_aesthetic_shadow, "model": self.aesthetic_shadow_model},
-            "waifu_scorer": {"name": "Waifu Scorer", "process": self._process_waifu_scorer, "model": self.waifu_scorer_model},
-            "aesthetic_predictor_v2_5": {"name": "Aesthetic V2.5", "process": self._process_aesthetic_predictor_v2_5, "model": self.aesthetic_predictor_model},
-            "anime_aesthetic": {"name": "Anime Score", "process": self._process_anime_aesthetic, "model": self.anime_aesthetic_model},
+    """Advanced model manager with async processing and resource management."""
+    
+    def __init__(self, processing_config: ProcessingConfig):
+        self.config = processing_config
+        self.models: Dict[ModelType, BaseModel] = {}
+        self.model_configs = self._create_default_configs()
+        self._processing_queue = asyncio.Queue()
+        self._worker_task: Optional[asyncio.Task] = None
+        self._temp_dir = Path(tempfile.mkdtemp())
+        
+        self._initialize_models()
+    
+    def _create_default_configs(self) -> Dict[ModelType, ModelConfig]:
+        """Create default model configurations."""
+        return {
+            ModelType.AESTHETIC_SHADOW: ModelConfig(
+                name="aesthetic_shadow",
+                display_name="Aesthetic Shadow"
+            ),
+            ModelType.WAIFU_SCORER: ModelConfig(
+                name="waifu_scorer",
+                display_name="Waifu Scorer"
+            ),
+            ModelType.AESTHETIC_PREDICTOR_V2_5: ModelConfig(
+                name="aesthetic_predictor_v2_5",
+                display_name="Aesthetic V2.5"
+            ),
+            ModelType.ANIME_AESTHETIC: ModelConfig(
+                name="anime_aesthetic",
+                display_name="Anime Score",
+                batch_supported=False
+            ),
         }
-        self.processing_queue: asyncio.Queue = asyncio.Queue()
-        self.worker_task = None # Initialize worker_task to None
-        self.temp_dir = tempfile.mkdtemp()
-
-    async def start_worker(self):
-        """Start the background worker task."""
-        if self.worker_task is None:
-            self.worker_task = asyncio.create_task(self._worker())
-
-    async def _worker(self):
-        """Background worker to process image evaluation requests from the queue."""
+    
+    def _initialize_models(self) -> None:
+        """Initialize all models."""
+        logger.info("Initializing models...")
+        
+        for model_type, config in self.model_configs.items():
+            if config.enabled:
+                try:
+                    model = ModelFactory.create_model(model_type, config, self.config.device)
+                    if model.is_available():
+                        self.models[model_type] = model
+                        logger.info(f"✓ {config.display_name} loaded successfully")
+                    else:
+                        logger.warning(f"✗ {config.display_name} failed to load")
+                except Exception as e:
+                    logger.error(f"✗ {config.display_name} initialization error: {e}")
+        
+        logger.info(f"Initialized {len(self.models)} models successfully")
+    
+    async def start_worker(self) -> None:
+        """Start the background processing worker."""
+        if self._worker_task is None:
+            self._worker_task = asyncio.create_task(self._worker_loop())
+            logger.info("Background worker started")
+    
+    async def _worker_loop(self) -> None:
+        """Main worker loop for processing requests."""
         while True:
-            request = await self.processing_queue.get()
-            if request is None: # Shutdown signal
-                self.processing_queue.task_done()
+            request = await self._processing_queue.get()
+            
+            if request is None:  # Shutdown signal
                 break
+            
             try:
-                results = await self._process_request(request)
-                request['results_future'].set_result(results) # Fulfill the future with results
+                result = await self._process_request(request)
+                request['future'].set_result(result)
             except Exception as e:
-                request['results_future'].set_exception(e) # Set exception if processing fails
+                request['future'].set_exception(e)
             finally:
-                self.processing_queue.task_done()
-
-    async def submit_request(self, request_data):
-        """Submit a new image processing request to the queue."""
-        results_future = asyncio.Future() # Future to hold the results
-        request = {**request_data, 'results_future': results_future}
-        await self.processing_queue.put(request)
-        return await results_future # Wait for and return results
-
-    async def _process_request(self, request):
-        """Process a single image evaluation request."""
+                self._processing_queue.task_done()
+    
+    async def process_images(
+        self, 
+        file_paths: List[str], 
+        selected_models: List[ModelType],
+        auto_batch: bool = False,
+        manual_batch_size: int = 1
+    ) -> BatchResult:
+        """Process images with selected models."""
+        future = asyncio.Future()
+        request = {
+            'file_paths': file_paths,
+            'selected_models': selected_models,
+            'auto_batch': auto_batch,
+            'manual_batch_size': manual_batch_size,
+            'future': future
+        }
+        
+        await self._processing_queue.put(request)
+        return await future
+    
+    async def _process_request(self, request: Dict) -> BatchResult:
+        """Process a single batch request."""
+        start_time = asyncio.get_event_loop().time()
+        logs = []
+        results = []
+        
         file_paths = request['file_paths']
+        selected_models = request['selected_models']
         auto_batch = request['auto_batch']
         manual_batch_size = request['manual_batch_size']
-        selected_models = request['selected_models']
-        log_events = []
+        
+        # Load images
+        images, valid_paths = await self._load_images(file_paths, logs)
+        
+        if not images:
+            return BatchResult([], logs, 0.0, 0, 0, len(file_paths))
+        
+        # Determine batch size
+        batch_size = await self._determine_batch_size(images, auto_batch, manual_batch_size, logs)
+        
+        # Process in batches
+        for i in range(0, len(images), batch_size):
+            batch_images = images[i:i+batch_size]
+            batch_paths = valid_paths[i:i+batch_size]
+            
+            batch_results = await self._process_batch(batch_images, batch_paths, selected_models, logs)
+            results.extend(batch_results)
+        
+        processing_time = asyncio.get_event_loop().time() - start_time
+        success_count = sum(1 for r in results if r.error is None)
+        error_count = len(results) - success_count
+        
+        return BatchResult(
+            results=results,
+            logs=logs,
+            processing_time=processing_time,
+            batch_size_used=batch_size,
+            success_count=success_count,
+            error_count=error_count
+        )
+    
+    async def _load_images(self, file_paths: List[str], logs: List[str]) -> Tuple[List[Image.Image], List[str]]:
+        """Load and validate images."""
         images = []
-        file_names = []
-        final_results = []
-
-        # Prepare images and file names
-        total_files = len(file_paths)
-        log_events.append(f"Starting to load {total_files} images...")
-        for f in file_paths:
+        valid_paths = []
+        
+        logs.append(f"Loading {len(file_paths)} images...")
+        
+        for path in file_paths:
             try:
-                img = Image.open(f).convert("RGB")
+                img = Image.open(path).convert("RGB")
                 images.append(img)
-                file_names.append(os.path.basename(f))
+                valid_paths.append(path)
             except Exception as e:
-                log_events.append(f"Error opening {f}: {e}")
-
-        if not images:
-            log_events.append("No valid images loaded.")
-            return [], log_events, 0, manual_batch_size
-
-        log_events.append("Images loaded. Determining batch size...")
-
-        try:
-            manual_batch_size = int(manual_batch_size) if manual_batch_size is not None else 1
-        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}")
-
-        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 = {}
-
-            # Process selected models
-            for model_key in selected_models:
-                if self.available_models[model_key]['selected']: # Ensure model is selected
-                    batch_results[model_key] = await self.available_models[model_key]['process'](batch_images, log_events) # Removed 'self' here
-                else:
-                    batch_results[model_key] = [None] * len(batch_images)
-
-            # Combine results and create final results list
-            for j in range(len(batch_images)):
-                scores_to_average = []
-                for model_key in selected_models:
-                     if self.available_models[model_key]['selected']: # Ensure model is selected
-                        score = batch_results[model_key][j]
-                        if score is not None:
-                            scores_to_average.append(score)
-
-                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 selected_models: # Add model scores to result
-                    if self.available_models[model_key]['selected']:
-                        result[model_key] = batch_results[model_key][j]
-                final_results.append(result)
-
-        log_events.append("All images processed.")
-        return final_results, log_events, 100, optimal_batch
-
-
-    def image_to_base64(self, image: Image.Image) -> str:
-        """Convert PIL Image to base64 encoded JPEG string."""
-        buffered = BytesIO()
-        image.save(buffered, format="JPEG")
-        return base64.b64encode(buffered.getvalue()).decode('utf-8')
-
-    def auto_tune_batch_size(self, images: list) -> int:
-        """Automatically determine the optimal batch size for processing."""
+                logs.append(f"Failed to load {path}: {e}")
+        
+        logs.append(f"Successfully loaded {len(images)} images")
+        return images, valid_paths
+    
+    async def _determine_batch_size(
+        self, 
+        images: List[Image.Image], 
+        auto_batch: bool, 
+        manual_batch_size: int,
+        logs: List[str]
+    ) -> int:
+        """Determine optimal batch size."""
+        if not auto_batch:
+            return min(manual_batch_size, len(images))
+        
+        # Auto-tune batch size
         batch_size = 1
-        max_batch = len(images)
         test_image = images[0:1]
-        while batch_size <= max_batch:
+        
+        while batch_size <= min(len(images), self.config.max_batch_size):
             try:
-                if "aesthetic_shadow" in self.available_models and self.available_models["aesthetic_shadow"]['selected']: # Check if model is available and selected
-                    _ = self.available_models["aesthetic_shadow"]['model'](test_image * batch_size)
-                if "waifu_scorer" in self.available_models and self.available_models["waifu_scorer"]['selected']: # Check if model is available and selected
-                    _ = self.available_models["waifu_scorer"]['model'](test_image * batch_size)
-                if "aesthetic_predictor_v2_5" in self.available_models and self.available_models["aesthetic_predictor_v2_5"]['selected']: # Check if model is available and selected
-                    _ = self.available_models["aesthetic_predictor_v2_5"]['model'].inference(test_image * batch_size)
+                # Test with a sample of available models
+                test_batch = test_image * batch_size
+                for model_type, model in list(self.models.items())[:2]:  # Test with first 2 models
+                    await model.predict(test_batch)
+                
                 batch_size *= 2
-                if batch_size > max_batch:
-                    break
             except Exception:
                 break
-        optimal = max(1, batch_size // 2)
-        if optimal > 64:
-            optimal = 64
-            print(f"Optimal batch size determined: {optimal}")
-        print(f"Optimal batch size determined: {optimal}")
-        return optimal
-
-    async def _process_aesthetic_shadow(self, batch_images, log_events):
-        try:
-            shadow_results = self.available_models["aesthetic_shadow"]['model'](batch_images)
-            log_events.append("Aesthetic Shadow processed for batch.")
-        except Exception as e:
-            log_events.append(f"Error in Aesthetic Shadow: {e}")
-            shadow_results = [None] * len(batch_images)
-        aesthetic_shadow_scores = []
-        for res in shadow_results:
-            try:
-                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.available_models["waifu_scorer"]['model'](batch_images)
-            waifu_scores = [float(np.clip(s, 0.0, 10.0)) if s is not None else None for s in waifu_scores]
-            log_events.append("Waifu Scorer processed for batch.")
-        except Exception as e:
-            log_events.append(f"Error in Waifu Scorer: {e}")
-            waifu_scores = [None] * len(batch_images)
-        return waifu_scores
-
-    async def _process_aesthetic_predictor_v2_5(self, batch_images, log_events):
-        try:
-            v2_5_scores = self.available_models["aesthetic_predictor_v2_5"]['model'].inference(batch_images)
-            v2_5_scores = [float(np.round(np.clip(s, 0.0, 10.0), 4)) if s is not None else None for s in v2_5_scores]
-            log_events.append("Aesthetic Predictor V2.5 processed for batch.")
-        except Exception as e:
-            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.available_models["anime_aesthetic"]['model'])
-                anime_scores.append(float(np.clip(score * 10.0, 0.0, 10.0)))
-                log_events.append(f"Anime Aesthetic processed for image {j + 1}.")
-            except Exception as e:
-                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 _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, results, 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"
+        
+        optimal_batch = max(1, batch_size // 2)
+        logs.append(f"Auto-tuned batch size: {optimal_batch}")
+        return optimal_batch
+    
+    async def _process_batch(
+        self, 
+        images: List[Image.Image], 
+        paths: List[str], 
+        selected_models: List[ModelType],
+        logs: List[str]
+    ) -> List[EvaluationResult]:
+        """Process a single batch of images."""
+        batch_results = []
+        
+        # Get predictions from all models
+        model_predictions = {}
+        for model_type in selected_models:
+            if model_type in self.models:
+                try:
+                    predictions = await self.models[model_type].predict(images)
+                    model_predictions[model_type.value] = predictions
+                    logs.append(f"✓ {self.model_configs[model_type].display_name} processed batch")
+                except Exception as e:
+                    logs.append(f"✗ {self.model_configs[model_type].display_name} error: {e}")
+                    model_predictions[model_type.value] = [None] * len(images)
+        
+        # Create results
+        for i, (image, path) in enumerate(zip(images, paths)):
+            # Collect scores for this image
+            scores = {}
+            valid_scores = []
+            
+            for model_type in selected_models:
+                score = model_predictions.get(model_type.value, [None] * len(images))[i]
+                scores[model_type.value] = score
+                if score is not None:
+                    valid_scores.append(score)
+            
+            # Calculate final score
+            final_score = np.mean(valid_scores) if valid_scores else None
+            if final_score is not None:
+                final_score = float(np.clip(final_score, *self.config.score_range))
+            
+            # Create thumbnail
+            thumbnail = image.copy()
+            thumbnail.thumbnail((200, 200), Image.Resampling.LANCZOS)
+            thumbnail_b64 = self._image_to_base64(thumbnail)
+            
+            result = EvaluationResult(
+                file_name=Path(path).name,
+                file_path=path,
+                thumbnail_b64=thumbnail_b64,
+                model_scores=scores,
+                final_score=final_score
+            )
+            
+            batch_results.append(result)
+        
+        return batch_results
+    
+    def _image_to_base64(self, image: Image.Image) -> str:
+        """Convert PIL Image to base64 string."""
+        buffer = BytesIO()
+        image.save(buffer, format="JPEG", quality=85, optimize=True)
+        return base64.b64encode(buffer.getvalue()).decode('utf-8')
+    
+    def get_available_models(self) -> Dict[ModelType, str]:
+        """Get available models with their display names."""
+        return {
+            model_type: self.model_configs[model_type].display_name 
+            for model_type in self.models.keys()
         }
-        key = key_map.get(sort_by, "final_score")
-        reverse = sort_by != "File Name"
-        results.sort(key=lambda r: r.get(key) if r.get(key) is not None else (-float('inf') if not reverse else float('inf')), reverse=reverse)
-        return results
-
-    def generate_html_table(self, results: list, selected_models) -> str:
-        """Generate an HTML table to display the evaluation results."""
-        table_html = """
+    
+    async def cleanup(self) -> None:
+        """Clean up resources."""
+        # Shutdown worker
+        if self._worker_task:
+            await self._processing_queue.put(None)
+            await self._worker_task
+        
+        # Clean up models
+        for model in self.models.values():
+            model.cleanup()
+        
+        # Clean up temp directory
+        if self._temp_dir.exists():
+            shutil.rmtree(self._temp_dir)
+        
+        logger.info("Model manager cleanup completed")
+
+
+# =============================================================================
+# Results Processing and Export
+# =============================================================================
+
+class ResultsProcessor:
+    """Handle result processing, sorting, and export functionality."""
+    
+    @staticmethod
+    def sort_results(results: List[EvaluationResult], sort_by: str, reverse: bool = True) -> List[EvaluationResult]:
+        """Sort results by specified criteria."""
+        sort_key_map = {
+            "Final Score": lambda r: r.final_score if r.final_score is not None else -float('inf'),
+            "File Name": lambda r: r.file_name.lower(),
+            **{f"model_{model_type.value}": lambda r, mt=model_type.value: r.model_scores.get(mt) or -float('inf') 
+               for model_type in ModelType}
+        }
+        
+        sort_key = sort_key_map.get(sort_by, sort_key_map["Final Score"])
+        return sorted(results, key=sort_key, reverse=reverse and sort_by != "File Name")
+    
+    @staticmethod
+    def generate_html_table(results: List[EvaluationResult], selected_models: List[ModelType]) -> str:
+        """Generate HTML table for results display."""
+        if not results:
+            return "<p>No results to display</p>"
+        
+        # CSS styles
+        styles = """
         <style>
-            .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; }
+            .results-table { 
+                width: 100%; border-collapse: collapse; margin: 20px 0; 
+                font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; 
+            }
+            .results-table th, .results-table td { 
+                border: 1px solid #ddd; padding: 12px; text-align: center; 
+            }
+            .results-table th { 
+                background-color: #f8f9fa; font-weight: 600; color: #495057; 
+            }
+            .results-table tr:nth-child(even) { background-color: #f8f9fa; }
+            .results-table tr:hover { background-color: #e9ecef; }
+            .image-preview { 
+                max-width: 120px; max-height: 120px; border-radius: 8px; 
+                box-shadow: 0 2px 4px rgba(0,0,0,0.1); 
+            }
+            .score-excellent { color: #28a745; font-weight: bold; }
+            .score-good { color: #ffc107; font-weight: bold; }
+            .score-poor { color: #dc3545; font-weight: bold; }
+            .score-na { color: #6c757d; font-style: italic; }
         </style>
-        <table class="results-table">
-            <thead>
-                <tr>
-                    <th>Image</th>
-                    <th>File Name</th>
         """
-        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>"
-
+        
+        # Table header
+        html = styles + '<table class="results-table"><thead><tr>'
+        html += '<th>Image</th><th>File Name</th>'
+        
+        for model_type in selected_models:
+            model_name = ModelType(model_type).name.replace('_', ' ').title()
+            html += f'<th>{model_name}</th>'
+        
+        html += '<th>Final Score</th></tr></thead><tbody>'
+        
+        # Table rows
         for result in results:
-            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>'
-
-
-    def cleanup(self):
-        """Clean up temporary directories and shutdown worker."""
-        if os.path.exists(self.temp_dir):
-            shutil.rmtree(self.temp_dir)
-        if self.worker_task is not None: # Check if worker_task was started
-            asyncio.run(self.shutdown()) # Shutdown worker gracefully
-
-    async def shutdown(self):
-        """Send shutdown signal to worker and wait for it to finish."""
-        if self.worker_task is not None: # Check if worker_task was started
-            await self.processing_queue.put(None) # Send shutdown signal
-            await self.worker_task # Wait for worker task to complete
-            await self.processing_queue.join() # Wait for queue to be empty
-
-
-#####################################
-#             Interface             #
-#####################################
-
-model_manager = ModelManager() # Initialize ModelManager once outside the interface function
-
-def create_interface():
-    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.
-
-        **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", 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_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")
-                download_file_output = gr.File() # Initialize gr.File component without filename
-                global_results_state = gr.State([]) # Initialize a global state to hold results
-
-        # Function to convert results to CSV format, excluding 'img_data'.
-        def results_to_csv(results, selected_models): # Take results as input
-            import csv
-            import io
-            if not 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 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, current_results):
-            file_paths = [f.name for f in files]
-
-            # Prepare request data for the ModelManager
-            request_data = {
-                'file_paths': file_paths,
-                'auto_batch': auto_batch,
-                'manual_batch_size': manual_batch,
-                'selected_models': {model: {'selected': model in selected_models} for model in model_options} # Pass model selections
-            }
-            # Submit request and get results from ModelManager
-            results, logs, progress_percent, updated_batch = await model_manager.submit_request(request_data)
-
-            updated_results = current_results + results # Append new results to current results
-
-            html_table = model_manager.generate_html_table(updated_results, selected_models)
-            progress_html = model_manager._generate_progress_html(progress_percent)
-            log_html = model_manager._format_logs(logs[-10:])
-
-            return status_html, html_table, log_html, progress_html, gr.update(value=updated_batch, interactive=not auto_batch), updated_results
-
-
-        def update_table_sort(sort_by_column, selected_models, current_results):
-            sorted_results = model_manager.sort_results(current_results, sort_by_column)
-            return model_manager.generate_html_table(sorted_results, selected_models), sorted_results # Return sorted results
-
-        def update_table_model_selection(selected_models, current_results):
-            # Recalculate final scores based on selected models
-            for result in current_results:
-                scores_to_average = []
-                for model_key in model_options: # Use model_options here, not available_models from manager in UI context
-                    if model_key in selected_models and model_key in model_manager.available_models and model_manager.available_models[model_key]['selected']: # consider only selected models from checkbox group and available_models
-                        score = result.get(model_key)
-                        if score is not None:
-                            scores_to_average.append(score)
-                final_score = float(np.clip(np.mean(scores_to_average), 0.0, 10.0)) if scores_to_average else None
-                result['final_score'] = final_score
-
-            sorted_results = model_manager.sort_results(current_results, "Final Score") # Keep sorting by Final Score when models change
-            return model_manager.generate_html_table(sorted_results, selected_models), sorted_results
-
-
-        def clear_results():
-            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),
-                    []) # Clear results state
-
-        def download_results_csv_trigger(selected_models, current_results): # Changed function name to avoid conflict and clarify purpose
-            csv_content = results_to_csv(current_results, selected_models)
-            if csv_content is None:
-                return None # Indicate no file to download
-
-            # 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
-
-
-        # Set initial selection state for models in ModelManager (important!)
-        for model_key in model_options:
-            model_manager.available_models[model_key]['selected'] = True # Default to all selected initially
-
-        auto_batch_checkbox.change(
-            update_batch_size_interactivity,
-            inputs=[auto_batch_checkbox],
-            outputs=[batch_size_input]
-        )
-
-        process_btn.click(
-            process_images_and_update,
-            inputs=[input_images, auto_batch_checkbox, batch_size_input, model_checkboxes, global_results_state],
-            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input, global_results_state]
-        )
-        sort_dropdown.change(
-            update_table_sort,
-            inputs=[sort_dropdown, model_checkboxes, global_results_state],
-            outputs=[output_html, global_results_state]
-        )
-        model_checkboxes.change( # Added change event for model checkboxes
-            update_table_model_selection,
-            inputs=[model_checkboxes, global_results_state],
-            outputs=[output_html, global_results_state]
-        )
-        clear_btn.click(
-            clear_results,
-            inputs=[],
-            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input, global_results_state]
-        )
-        download_csv.click(
-            download_results_csv_trigger, # Call the trigger function
-            inputs=[model_checkboxes, global_results_state],
-            outputs=[download_file_output] # Output is now the gr.File component
-        )
-        demo.load(lambda: update_table_sort("Final Score", model_options, []), inputs=None, outputs=[output_html, global_results_state]) # Initial sort and table render, pass empty initial results
-        demo.load(model_manager.start_worker) # Start the worker task on demo load
-
-        gr.Markdown("""
-        ### Notes
-        - 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
-
-if __name__ == "__main__":
-    demo = create_interface()
-    demo.queue().launch()
+            html += '<tr>'
+            html += f'<td><img src="data:image/jpeg;base64,{result.thumbnail_b64}" class="image-preview" alt="{result.file_name}"></td>'
+            html += f'<td>{result.file_name}</td>'
+            
+            # Model scores
+            for model_type in selected_models:
+                score = result.model_scores.get(model_type.value)
+                html += ResultsProcessor._format_score_cell(score)
+            
+            # Final score
+            html += ResultsProcessor._format_score_cell(result.final_score)
+            html += '</tr>'
+        
+        html += '</tbody></table>'
+        return html