Update app.py

app.py

@@ -1,824 +1,600 @@
-"""
-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 tempfile
-import shutil
-from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
-from enum import Enum
-from io import BytesIO, StringIO
+import base64
+from io import BytesIO
+from typing import List, Dict, Any, Optional, Tuple
+from dataclasses import dataclass
 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 torch.nn as nn
+import onnxruntime as rt
 from PIL import Image
+import gradio as gr
+import pandas as pd
 from transformers import pipeline
 from huggingface_hub import hf_hub_download
 
-# 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)
+# Import necessary function from aesthetic_predictor_v2_5
+from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip
 
 
 @dataclass
 class EvaluationResult:
-    """Data class for individual evaluation results."""
+    """Data class for storing image evaluation results."""
     file_name: str
-    file_path: str
-    thumbnail_b64: str
-    model_scores: Dict[str, Optional[float]] = field(default_factory=dict)
+    image: Image.Image
+    aesthetic_shadow: Optional[float] = None
+    waifu_scorer: Optional[float] = None
+    aesthetic_v2_5: Optional[float] = None
+    anime_aesthetic: Optional[float] = None
     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
-    ):
+class MLP(torch.nn.Module):
+    """Optimized MLP for image feature regression."""
+    def __init__(self, input_size: int = 768):
         super().__init__()
-        
-        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)
-    
+        self.network = torch.nn.Sequential(
+            torch.nn.Linear(input_size, 1024),
+            torch.nn.ReLU(),
+            torch.nn.BatchNorm1d(1024),
+            torch.nn.Dropout(0.2),
+            torch.nn.Linear(1024, 256),
+            torch.nn.ReLU(),
+            torch.nn.BatchNorm1d(256),
+            torch.nn.Dropout(0.1),
+            torch.nn.Linear(256, 64),
+            torch.nn.ReLU(),
+            torch.nn.Linear(64, 1)
+        )
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.network(x)
 
 
-class WaifuScorerModel:
-    """Modernized WaifuScorer implementation with better error handling."""
+class ModelLoader:
+    """Centralized model loading and management."""
     
-    def __init__(self, config: ModelConfig, device: str):
-        self.config = config
-        self.device = device
-        self.dtype = torch.float32
-        self._available = False
-        self._model = None
-        self._clip_model = None
-        self._preprocess = None
-        
-        self._initialize_model()
+    def __init__(self, device: str = None):
+        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        self.models = {}
+        self._load_all_models()
     
-    def _initialize_model(self) -> None:
-        """Initialize the model with proper error handling."""
+    def _load_all_models(self):
+        """Load all models during initialization."""
+        try:
+            self._load_aesthetic_shadow()
+            self._load_waifu_scorer()
+            self._load_aesthetic_v2_5()
+            self._load_anime_aesthetic()
+            print("✅ All models loaded successfully!")
+        except Exception as e:
+            print(f"❌ Error loading models: {e}")
+    
+    def _load_aesthetic_shadow(self):
+        """Load Aesthetic Shadow model."""
+        print("🔄 Loading Aesthetic Shadow...")
+        self.models['aesthetic_shadow'] = pipeline(
+            "image-classification",
+            model="NeoChen1024/aesthetic-shadow-v2-backup",
+            device=self.device
+        )
+    
+    def _load_waifu_scorer(self):
+        """Load Waifu Scorer model."""
+        print("🔄 Loading Waifu Scorer...")
         try:
             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=self.device)
+            # Load MLP
+            model_path = hf_hub_download("Eugeoter/waifu-scorer-v3", "model.pth")
+            mlp = MLP()
+            state_dict = torch.load(model_path, map_location=self.device)
+            mlp.load_state_dict(state_dict)
+            mlp.to(self.device).eval()
             
-            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}")
+            # Load CLIP
+            clip_model, preprocess = clip.load("ViT-L/14", device=self.device)
             
+            self.models['waifu_scorer'] = {
+                'mlp': mlp,
+                'clip_model': clip_model,
+                'preprocess': preprocess
+            }
         except Exception as e:
-            logger.error(f"Failed to initialize WaifuScorer: {e}")
-            self._available = False
+            print(f"⚠️ Waifu Scorer not available: {e}")
+            self.models['waifu_scorer'] = None
     
-    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)
-        
+    def _load_aesthetic_v2_5(self):
+        """Load Aesthetic Predictor V2.5."""
+        print("🔄 Loading Aesthetic V2.5...")
         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)]
+            model, preprocessor = convert_v2_5_from_siglip(
+                low_cpu_mem_usage=True,
+                trust_remote_code=True,
+            )
+            if torch.cuda.is_available():
+                model = model.to(torch.bfloat16).cuda()
             
+            self.models['aesthetic_v2_5'] = {
+                'model': model,
+                'preprocessor': preprocessor
+            }
         except Exception as e:
-            logger.error(f"Error in WaifuScorer prediction: {e}")
-            return [None] * len(images)
-    
-    def is_available(self) -> bool:
-        return self._available
+            print(f"⚠️ Aesthetic V2.5 not available: {e}")
+            self.models['aesthetic_v2_5'] = None
     
-    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
+    def _load_anime_aesthetic(self):
+        """Load Anime Aesthetic model."""
+        print("🔄 Loading Anime Aesthetic...")
+        try:
+            model_path = hf_hub_download("skytnt/anime-aesthetic", "model.onnx")
+            self.models['anime_aesthetic'] = rt.InferenceSession(
+                model_path, 
+                providers=['CPUExecutionProvider']
+            )
+        except Exception as e:
+            print(f"⚠️ Anime Aesthetic not available: {e}")
+            self.models['anime_aesthetic'] = None
 
 
-class AestheticShadowModel:
-    """Wrapper for Aesthetic Shadow model using transformers pipeline."""
+class ImageEvaluator:
+    """Main image evaluation class with batch processing."""
     
-    def __init__(self, config: ModelConfig, device: str):
-        self.config = config
-        self.device = device
-        self._available = False
-        self._model = None
-        
-        self._initialize_model()
+    def __init__(self):
+        self.loader = ModelLoader()
+        self.temp_dir = Path(tempfile.mkdtemp())
     
-    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")
+    def evaluate_images(
+        self,
+        images: List[Image.Image],
+        file_names: List[str],
+        selected_models: List[str],
+        batch_size: int = 4,
+        progress_callback=None
+    ) -> List[EvaluationResult]:
+        """Evaluate images using selected models."""
+        results = []
+        total_batches = (len(images) + batch_size - 1) // batch_size
+        
+        for batch_idx in range(0, len(images), batch_size):
+            batch_images = images[batch_idx:batch_idx + batch_size]
+            batch_names = file_names[batch_idx:batch_idx + batch_size]
             
-        except Exception as e:
-            logger.error(f"Failed to initialize Aesthetic Shadow: {e}")
-            self._available = False
+            # Update progress
+            if progress_callback:
+                progress = (batch_idx // batch_size + 1) / total_batches
+                progress_callback(progress, f"Processing batch {batch_idx//batch_size + 1}/{total_batches}")
+            
+            # Process batch
+            batch_results = self._process_batch(batch_images, batch_names, selected_models)
+            results.extend(batch_results)
+        
+        return results
     
-    async def predict(self, images: List[Image.Image]) -> List[Optional[float]]:
-        """Predict scores for a batch of images."""
-        if not self._available:
+    def _process_batch(
+        self,
+        images: List[Image.Image],
+        file_names: List[str],
+        selected_models: List[str]
+    ) -> List[EvaluationResult]:
+        """Process a single batch of images."""
+        batch_results = []
+        
+        # Initialize results
+        for i, (img, name) in enumerate(zip(images, file_names)):
+            result = EvaluationResult(file_name=name, image=img)
+            batch_results.append(result)
+        
+        # Process each selected model
+        if 'aesthetic_shadow' in selected_models:
+            scores = self._eval_aesthetic_shadow(images)
+            for result, score in zip(batch_results, scores):
+                result.aesthetic_shadow = score
+        
+        if 'waifu_scorer' in selected_models:
+            scores = self._eval_waifu_scorer(images)
+            for result, score in zip(batch_results, scores):
+                result.waifu_scorer = score
+        
+        if 'aesthetic_v2_5' in selected_models:
+            scores = self._eval_aesthetic_v2_5(images)
+            for result, score in zip(batch_results, scores):
+                result.aesthetic_v2_5 = score
+        
+        if 'anime_aesthetic' in selected_models:
+            scores = self._eval_anime_aesthetic(images)
+            for result, score in zip(batch_results, scores):
+                result.anime_aesthetic = score
+        
+        # Calculate final scores
+        for result in batch_results:
+            result.final_score = self._calculate_final_score(result, selected_models)
+        
+        return batch_results
+    
+    def _eval_aesthetic_shadow(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Evaluate using Aesthetic Shadow model."""
+        if not self.loader.models.get('aesthetic_shadow'):
             return [None] * len(images)
         
         try:
-            results = self._model(images)
+            results = self.loader.models['aesthetic_shadow'](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(float(np.clip(hq_score * 10.0, 0.0, 10.0)))
+                except:
                     scores.append(None)
-            
             return scores
-            
         except Exception as e:
-            logger.error(f"Error in Aesthetic Shadow prediction: {e}")
+            print(f"Error in Aesthetic Shadow: {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."""
+    def _eval_waifu_scorer(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Evaluate using Waifu Scorer model."""
+        model_dict = self.loader.models.get('waifu_scorer')
+        if not model_dict:
+            return [None] * len(images)
+        
         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()
-            
-            self._available = True
-            logger.info("Aesthetic Predictor V2.5 loaded successfully")
-            
+            with torch.no_grad():
+                # Preprocess images
+                image_tensors = [model_dict['preprocess'](img).unsqueeze(0) for img in images]
+                if len(image_tensors) == 1:
+                    image_tensors = image_tensors * 2  # CLIP requirement
+                
+                image_batch = torch.cat(image_tensors).to(self.loader.device)
+                image_features = model_dict['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(self.loader.device)
+                
+                predictions = model_dict['mlp'](im_emb)
+                scores = predictions.clamp(0, 10).cpu().numpy().flatten().tolist()
+                
+                return scores[:len(images)]
         except Exception as e:
-            logger.error(f"Failed to initialize Aesthetic Predictor V2.5: {e}")
-            self._available = False
+            print(f"Error in Waifu Scorer: {e}")
+            return [None] * len(images)
     
-    async def predict(self, images: List[Image.Image]) -> List[Optional[float]]:
-        """Predict scores for a batch of images."""
-        if not self._available:
+    def _eval_aesthetic_v2_5(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Evaluate using Aesthetic Predictor V2.5."""
+        model_dict = self.loader.models.get('aesthetic_v2_5')
+        if not model_dict:
             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
+            pixel_values = model_dict['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]
-            
+                scores = model_dict['model'](pixel_values).logits.squeeze().float().cpu().numpy()
+                if scores.ndim == 0:
+                    scores = np.array([scores])
+                
+                return [float(np.clip(s, 0.0, 10.0)) for s in scores.tolist()]
         except Exception as e:
-            logger.error(f"Error in Aesthetic Predictor V2.5 prediction: {e}")
+            print(f"Error in Aesthetic V2.5: {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:
+    def _eval_anime_aesthetic(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Evaluate using Anime Aesthetic model."""
+        model = self.loader.models.get('anime_aesthetic')
+        if not model:
             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)))
+                # Preprocess image
+                img_np = np.array(img).astype(np.float32) / 255.0
+                h, w = img_np.shape[:2]
+                s = 768
+                
+                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))
+                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
+                
+                input_tensor = np.transpose(canvas, (2, 0, 1))[np.newaxis, :]
+                pred = model.run(None, {"img": input_tensor})[0].item()
+                scores.append(float(np.clip(pred * 10.0, 0.0, 10.0)))
             except Exception as e:
-                logger.error(f"Error predicting anime aesthetic for image: {e}")
+                print(f"Error processing 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:
-    """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
-            ),
-        }
-    
-    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
-                break
-            
-            try:
-                result = await self._process_request(request)
-                request['future'].set_result(result)
-            except Exception as e:
-                request['future'].set_exception(e)
-            finally:
-                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']
-        
-        # 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 = []
-        valid_paths = []
+    def _calculate_final_score(self, result: EvaluationResult, selected_models: List[str]) -> Optional[float]:
+        """Calculate final score from selected model results."""
+        scores = []
         
-        logs.append(f"Loading {len(file_paths)} images...")
+        for model in selected_models:
+            score = getattr(result, model, None)
+            if score is not None:
+                scores.append(score)
         
-        for path in file_paths:
-            try:
-                img = Image.open(path).convert("RGB")
-                images.append(img)
-                valid_paths.append(path)
-            except Exception as e:
-                logs.append(f"Failed to load {path}: {e}")
-        
-        logs.append(f"Successfully loaded {len(images)} images")
-        return images, valid_paths
+        return float(np.mean(scores)) if scores else None
     
-    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
+    def results_to_dataframe(self, results: List[EvaluationResult]) -> pd.DataFrame:
+        """Convert results to pandas DataFrame."""
+        data = []
+        for result in results:
+            row = {
+                'File Name': result.file_name,
+                'Final Score': result.final_score,
+            }
+            if result.aesthetic_shadow is not None:
+                row['Aesthetic Shadow'] = result.aesthetic_shadow
+            if result.waifu_scorer is not None:
+                row['Waifu Scorer'] = result.waifu_scorer
+            if result.aesthetic_v2_5 is not None:
+                row['Aesthetic V2.5'] = result.aesthetic_v2_5
+            if result.anime_aesthetic is not None:
+                row['Anime Aesthetic'] = result.anime_aesthetic
+            data.append(row)
+        
+        return pd.DataFrame(data)
+    
+    def optimize_batch_size(self, sample_images: List[Image.Image]) -> int:
+        """Automatically determine optimal batch size."""
+        if not sample_images:
+            return 1
+        
+        test_image = sample_images[0]
         batch_size = 1
-        test_image = images[0:1]
+        max_test = min(16, len(sample_images))
         
-        while batch_size <= min(len(images), self.config.max_batch_size):
+        while batch_size <= max_test:
             try:
-                # 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)
-                
+                test_batch = [test_image] * batch_size
+                # Test with a lightweight model
+                if self.loader.models.get('aesthetic_shadow'):
+                    _ = self.loader.models['aesthetic_shadow'](test_batch)
                 batch_size *= 2
             except Exception:
                 break
         
-        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 = []
+        optimal = max(1, batch_size // 2)
+        return min(optimal, 8)  # Cap at reasonable size
+
+
+def create_interface():
+    """Create the Gradio interface."""
+    evaluator = ImageEvaluator()
+    
+    # Available models
+    model_choices = [
+        ("Aesthetic Shadow", "aesthetic_shadow"),
+        ("Waifu Scorer", "waifu_scorer"),
+        ("Aesthetic V2.5", "aesthetic_v2_5"),
+        ("Anime Aesthetic", "anime_aesthetic")
+    ]
+    available_models = [choice[1] for choice in model_choices]
+    
+    with gr.Blocks(title="Image Evaluation Tool", theme=gr.themes.Soft()) as app:
+        gr.Markdown("""
+        # 🎨 Modern Image Evaluation Tool
+        
+        Upload images to evaluate them using state-of-the-art aesthetic and quality prediction models.
+        
+        **Features:**
+        - Multiple AI models for comprehensive evaluation
+        - Batch processing with automatic optimization
+        - Interactive results table with sorting and filtering
+        - CSV export functionality
+        - Real-time progress tracking
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Input components
+                input_files = gr.File(
+                    label="📁 Upload Images",
+                    file_count="multiple",
+                    file_types=["image"]
+                )
+                
+                model_selection = gr.CheckboxGroup(
+                    choices=model_choices,
+                    value=available_models,
+                    label="🤖 Select Models",
+                    info="Choose which models to use for evaluation"
+                )
+                
+                with gr.Row():
+                    auto_batch = gr.Checkbox(
+                        label="🔄 Auto Batch Size",
+                        value=True,
+                        info="Automatically optimize batch size"
+                    )
+                    
+                    manual_batch = gr.Slider(
+                        minimum=1,
+                        maximum=16,
+                        value=4,
+                        step=1,
+                        label="📊 Batch Size",
+                        interactive=False,
+                        info="Manual batch size (when auto is disabled)"
+                    )
+                
+                evaluate_btn = gr.Button(
+                    "🚀 Evaluate Images",
+                    variant="primary",
+                    size="lg"
+                )
+                
+                clear_btn = gr.Button("🗑️ Clear Results", variant="secondary")
+            
+            with gr.Column(scale=2):
+                # Progress and status
+                progress_bar = gr.Progress()
+                status_text = gr.Textbox(
+                    label="📊 Status",
+                    interactive=False,
+                    max_lines=2
+                )
+                
+                # Results display
+                results_table = gr.DataFrame(
+                    label="📋 Evaluation Results",
+                    interactive=False,
+                    wrap=True,
+                    max_height=400
+                )
+                
+                # Export functionality
+                with gr.Row():
+                    export_csv = gr.Button("📥 Export CSV", variant="secondary")
+                    download_file = gr.File(
+                        label="💾 Download",
+                        visible=False
+                    )
         
-        # 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 = []
+        # State management
+        results_state = gr.State([])
+        
+        # Event handlers
+        def toggle_batch_slider(auto_enabled):
+            return gr.update(interactive=not auto_enabled)
+        
+        def process_images(files, models, auto_batch_enabled, manual_batch_size, progress=gr.Progress()):
+            if not files or not models:
+                return "❌ Please upload images and select at least one model", pd.DataFrame(), []
             
-            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)
+            try:
+                # Load images
+                images = []
+                file_names = []
+                
+                progress(0.1, "📂 Loading images...")
+                
+                for file in files:
+                    try:
+                        img = Image.open(file.name).convert("RGB")
+                        images.append(img)
+                        file_names.append(os.path.basename(file.name))
+                    except Exception as e:
+                        print(f"Error loading {file.name}: {e}")
+                
+                if not images:
+                    return "❌ No valid images loaded", pd.DataFrame(), []
+                
+                # Determine batch size
+                if auto_batch_enabled:
+                    batch_size = evaluator.optimize_batch_size(images[:2])
+                    progress(0.2, f"🔧 Optimized batch size: {batch_size}")
+                else:
+                    batch_size = int(manual_batch_size)
+                
+                # Process images
+                def progress_callback(prog, msg):
+                    progress(0.2 + prog * 0.7, msg)
+                
+                results = evaluator.evaluate_images(
+                    images, file_names, models, batch_size, progress_callback
+                )
+                
+                progress(0.95, "📊 Generating results table...")
+                
+                # Convert to DataFrame
+                df = evaluator.results_to_dataframe(results)
+                df = df.sort_values('Final Score', ascending=False, na_position='last')
+                
+                progress(1.0, f"✅ Processed {len(results)} images successfully!")
+                
+                return f"✅ Evaluated {len(results)} images using {len(models)} models", df, results
+                
+            except Exception as e:
+                return f"❌ Error during processing: {str(e)}", pd.DataFrame(), []
+        
+        def update_results_table(models, current_results):
+            if not current_results:
+                return pd.DataFrame()
             
-            # 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))
+            # Recalculate final scores based on selected models
+            for result in current_results:
+                result.final_score = evaluator._calculate_final_score(result, models)
             
-            # Create thumbnail
-            thumbnail = image.copy()
-            thumbnail.thumbnail((200, 200), Image.Resampling.LANCZOS)
-            thumbnail_b64 = self._image_to_base64(thumbnail)
+            df = evaluator.results_to_dataframe(current_results)
+            return df.sort_values('Final Score', ascending=False, na_position='last')
+        
+        def export_results(current_results):
+            if not current_results:
+                return gr.update(visible=False)
             
-            result = EvaluationResult(
-                file_name=Path(path).name,
-                file_path=path,
-                thumbnail_b64=thumbnail_b64,
-                model_scores=scores,
-                final_score=final_score
-            )
+            df = evaluator.results_to_dataframe(current_results)
+            csv_path = evaluator.temp_dir / "evaluation_results.csv"
+            df.to_csv(csv_path, index=False)
             
-            batch_results.append(result)
+            return gr.update(value=str(csv_path), visible=True)
         
-        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()
-        }
-    
-    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: '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>
-        """
+        def clear_all():
+            return (
+                "🔄 Ready for new evaluation",
+                pd.DataFrame(),
+                [],
+                gr.update(visible=False)
+            )
         
-        # Table header
-        html = styles + '<table class="results-table"><thead><tr>'
-        html += '<th>Image</th><th>File Name</th>'
+        # Wire up events
+        auto_batch.change(
+            toggle_batch_slider,
+            inputs=[auto_batch],
+            outputs=[manual_batch]
+        )
         
-        for model_type in selected_models:
-            model_name = ModelType(model_type).name.replace('_', ' ').title()
-            html += f'<th>{model_name}</th>'
+        evaluate_btn.click(
+            process_images,
+            inputs=[input_files, model_selection, auto_batch, manual_batch],
+            outputs=[status_text, results_table, results_state]
+        )
         
-        html += '<th>Final Score</th></tr></thead><tbody>'
+        model_selection.change(
+            update_results_table,
+            inputs=[model_selection, results_state],
+            outputs=[results_table]
+        )
         
-        # Table rows
-        for result in results:
-            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>'
+        export_csv.click(
+            export_results,
+            inputs=[results_state],
+            outputs=[download_file]
+        )
+        
+        clear_btn.click(
+            clear_all,
+            outputs=[status_text, results_table, results_state, download_file]
+        )
         
-        html += '</tbody></table>'
-        return html
+        # Initial setup
+        app.load(lambda: "🔄 Ready for evaluation - Upload images to get started!")
+    
+    return app
+
+
+if __name__ == "__main__":
+    app = create_interface()
+    app.queue(max_size=10).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True
+    )