Update app.py

app.py

@@ -12,15 +12,16 @@ import torch
 import onnxruntime as rt
 from PIL import Image
 from huggingface_hub import hf_hub_download
-from transformers import pipeline, Pipeline
+from transformers import pipeline, Pipeline, AutoModel, AutoProcessor
 from tqdm import tqdm
 
-# Suppress a specific PIL warning about image size
+# Suppress a specific PIL warning about image size to handle large images
 Image.MAX_IMAGE_PIXELS = None
 
 # --- Configuration ---
 CACHE_DIR = "./hf_cache"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# Use bfloat16 for modern GPUs, float32 for others (including CPU)
 DTYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 else torch.float32
 
 print(f"Using device: {DEVICE} with dtype: {DTYPE}")
@@ -31,21 +32,20 @@ print(f"Using device: {DEVICE} with dtype: {DTYPE}")
 
 class AestheticScorer(ABC):
     """Abstract base class for all aesthetic scoring models."""
-
     def __init__(self, model_name: str, repo_id: str, filename: str = None):
         self.model_name = model_name
         self.repo_id = repo_id
         self.filename = filename
         self._model = None
-        print(f"Initializing scorer: {self.model_name}")
+        print(f"Initializing scorer definition: {self.model_name}")
 
     @property
     def model(self):
         """Lazy-loads the model on first access."""
         if self._model is None:
-            print(f"Loading model for '{self.model_name}'...")
+            print(f"Loading model weights for '{self.model_name}'...")
             self._model = self.load_model()
-            print(f"'{self.model_name}' model loaded.")
+            print(f"'{self.model_name}' model weights loaded.")
         return self._model
 
     def _download_model(self) -> str:
@@ -63,9 +63,9 @@ class AestheticScorer(ABC):
         pass
 
     def release_model(self):
-        """Releases model from memory."""
+        """Releases model from memory to conserve VRAM/RAM."""
         if self._model is not None:
-            print(f"Releasing model: {self.model_name}")
+            print(f"Releasing model from memory: {self.model_name}")
             del self._model
             self._model = None
             gc.collect()
@@ -74,23 +74,15 @@ class AestheticScorer(ABC):
 
 class PipelineScorer(AestheticScorer):
     """Scorer for models compatible with Hugging Face pipelines."""
-
     def load_model(self) -> Pipeline:
-        """Loads a pipeline model."""
-        return pipeline(
-            "image-classification",
-            model=self.repo_id,
-            device=DEVICE,
-        )
+        return pipeline("image-classification", model=self.repo_id, device=DEVICE)
 
     @torch.no_grad()
     def score_batch(self, image_batch: List[Image.Image]) -> List[float]:
-        """Scores a batch using the pipeline and extracts the 'hq' score."""
-        results = self.model(image_batch)
+        results = self.model(image_batch, top_k=None) # Get all class scores
         scores = []
         for res in results:
             try:
-                # Find the score for the 'hq' (high quality) label
                 hq_score = next(item['score'] for item in res if item['label'] == 'hq')
                 scores.append(round(hq_score * 10.0, 4))
             except (StopIteration, TypeError):
@@ -99,21 +91,16 @@ class PipelineScorer(AestheticScorer):
 
 class ONNXScorer(AestheticScorer):
     """Scorer for ONNX-based models."""
-
     def load_model(self) -> rt.InferenceSession:
-        """Loads an ONNX inference session."""
         model_path = self._download_model()
         return rt.InferenceSession(model_path, providers=['CUDAExecutionProvider' if DEVICE == 'cuda' else 'CPUExecutionProvider'])
 
     def _preprocess(self, img: Image.Image) -> np.ndarray:
-        """Preprocesses a single image for the Anime Aesthetic model."""
         img_np = np.array(img.convert("RGB")).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
+        ratio = s / max(h, w)
+        new_h, new_w = int(h * ratio), int(w * ratio)
         
         resized = cv2.resize(img_np, (new_w, new_h), interpolation=cv2.INTER_AREA)
         canvas = np.zeros((s, s, 3), dtype=np.float32)
@@ -123,7 +110,6 @@ class ONNXScorer(AestheticScorer):
         return np.transpose(canvas, (2, 0, 1))[np.newaxis, :]
 
     def score_batch(self, image_batch: List[Image.Image]) -> List[float]:
-        """Scores images one by one as this model doesn't support batching."""
         scores = []
         for img in image_batch:
             try:
@@ -135,264 +121,192 @@ class ONNXScorer(AestheticScorer):
         return scores
 
 class CLIPMLPScorer(AestheticScorer):
-    """Scorer for models using a CLIP backbone and an MLP head."""
-
+    """Scorer for models using a CLIP backbone and a custom MLP head."""
     class MLP(torch.nn.Module):
+        """Re-implementation of the exact MLP from the original code."""
         def __init__(self, input_size: int):
             super().__init__()
             self.layers = torch.nn.Sequential(
-                torch.nn.Linear(input_size, 1024),
+                torch.nn.Linear(input_size, 2048),
                 torch.nn.ReLU(),
-                torch.nn.Dropout(0.2),
-                torch.nn.Linear(1024, 128),
+                torch.nn.BatchNorm1d(2048),
+                torch.nn.Dropout(0.3),
+                torch.nn.Linear(2048, 512),
+                torch.nn.ReLU(),
+                torch.nn.BatchNorm1d(512),
+                torch.nn.Dropout(0.3),
+                torch.nn.Linear(512, 256),
                 torch.nn.ReLU(),
+                torch.nn.BatchNorm1d(256),
                 torch.nn.Dropout(0.2),
-                torch.nn.Linear(128, 64),
+                torch.nn.Linear(256, 128),
                 torch.nn.ReLU(),
-                torch.nn.Linear(64, 16),
+                torch.nn.BatchNorm1d(128),
+                torch.nn.Dropout(0.1),
+                torch.nn.Linear(128, 32),
                 torch.nn.ReLU(),
-                torch.nn.Linear(16, 1),
+                torch.nn.Linear(32, 1)
             )
         def forward(self, x):
             return self.layers(x)
 
     def load_model(self) -> Dict[str, Any]:
-        """Loads both the CLIP model and the custom MLP head."""
-        import clip # Lazy import
-        
+        import clip
         model_path = self._download_model()
-        
         mlp = self.MLP(input_size=768) # ViT-L/14 has 768 features
         state_dict = torch.load(model_path, map_location=DEVICE)
         mlp.load_state_dict(state_dict)
-        mlp.to(device=DEVICE, dtype=DTYPE)
+        mlp.to(device=DEVICE)
         mlp.eval()
-
         clip_model, preprocess = clip.load("ViT-L/14", device=DEVICE)
-        
         return {"mlp": mlp, "clip": clip_model, "preprocess": preprocess}
 
     @torch.no_grad()
     def score_batch(self, image_batch: List[Image.Image]) -> List[float]:
-        """Scores a batch using CLIP features and the MLP head."""
         preprocess = self.model['preprocess']
+        # Handle single-image batches correctly for CLIP
+        if len(image_batch) == 1:
+            image_batch = image_batch * 2
+            single_image_mode = True
+        else:
+            single_image_mode = False
+
         image_tensors = torch.cat([preprocess(img).unsqueeze(0) for img in image_batch]).to(DEVICE)
-        
-        image_features = self.model['clip'].encode_image(image_tensors)
+        image_features = self.model['clip'].encode_image(image_tensors).to(torch.float32)
         image_features /= image_features.norm(dim=-1, keepdim=True)
+        predictions = self.model['mlp'](image_features).squeeze(-1)
+        scores = predictions.clamp(0, 10).float().cpu().numpy()
         
-        # Pass features through MLP
-        predictions = self.model['mlp'](image_features.to(DTYPE)).squeeze(-1)
-        scores = predictions.float().cpu().numpy()
-        
+        final_scores = [round(float(s), 4) for s in scores]
+        return final_scores[:1] if single_image_mode else final_scores
+
+class SigLIPScorer(AestheticScorer):
+    """Scorer for the Aesthetic Predictor V2.5 SigLIP model."""
+    def load_model(self) -> Dict[str, Any]:
+        model = AutoModel.from_pretrained(self.repo_id, trust_remote_code=True).to(DEVICE, DTYPE).eval()
+        processor = AutoProcessor.from_pretrained(self.repo_id, trust_remote_code=True)
+        return {"model": model, "processor": processor}
+
+    @torch.no_grad()
+    def score_batch(self, image_batch: List[Image.Image]) -> List[float]:
+        inputs = self.model['processor'](
+            images=[img.convert("RGB") for img in image_batch],
+            return_tensors="pt"
+        )
+        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+        inputs['pixel_values'] = inputs['pixel_values'].to(DTYPE)
+        logits = self.model(**inputs).logits.squeeze(-1)
+        scores = logits.float().cpu().numpy()
         return [round(float(s), 4) for s in scores]
 
 # --- Model Registry ---
-MODEL_REGISTRY: Dict[str, Type[AestheticScorer]] = {
-    "Aesthetic Shadow V2": PipelineScorer(
-        "Aesthetic Shadow V2", "shadowlilac/aesthetic-shadow-v2"
-    ),
-    "Waifu Scorer V2": CLIPMLPScorer(
-        "Waifu Scorer V2", "skytnt/waifu-aesthetic-scorer", "model.pth"
-    ),
-    "Anime Scorer": ONNXScorer(
-        "Anime Scorer", "skytnt/anime-aesthetic", "model.onnx"
-    )
+MODEL_REGISTRY: Dict[str, AestheticScorer] = {
+    "Aesthetic Shadow V2": PipelineScorer("Aesthetic Shadow V2", "NeoChen1024/aesthetic-shadow-v2-backup"),
+    "Waifu Scorer V3": CLIPMLPScorer("Waifu Scorer V3", "Eugeoter/waifu-scorer-v3", "model.pth"),
+    "Aesthetic V2.5 SigLIP": SigLIPScorer("Aesthetic V2.5 SigLIP", "জিংוניत्र/Aesthetic-Predictor-V2-5-SigLIP"),
+    "Anime Scorer": ONNXScorer("Anime Scorer", "skytnt/anime-aesthetic", "model.onnx")
 }
-
-# In-memory cache for loaded model instances
 _loaded_models_cache: Dict[str, AestheticScorer] = {}
 
-
 # ==================================================================================
 # 2. CORE PROCESSING LOGIC
 # ==================================================================================
 
 def get_scorers(model_names: List[str]) -> List[AestheticScorer]:
     """Retrieves and caches scorer instances based on selected names."""
-    # Release models that are no longer selected
-    for name, scorer in list(_loaded_models_cache.items()):
+    for name in list(_loaded_models_cache.keys()):
         if name not in model_names:
-            scorer.release_model()
+            _loaded_models_cache[name].release_model()
             del _loaded_models_cache[name]
-
-    # Load newly selected models
-    scorers = []
-    for name in model_names:
-        if name in _loaded_models_cache:
-            scorers.append(_loaded_models_cache[name])
-        elif name in MODEL_REGISTRY:
-            scorer = MODEL_REGISTRY[name]
-            _loaded_models_cache[name] = scorer
-            scorers.append(scorer)
-    return scorers
+    return [_loaded_models_cache.setdefault(name, MODEL_REGISTRY[name]) for name in model_names]
     
 def evaluate_images(
-    files: List[gr.File],
-    selected_model_names: List[str],
-    batch_size: int,
-    progress: gr.Progress = gr.Progress(track_tqdm=True),
+    files: List[gr.File], selected_model_names: List[str], batch_size: int, progress=gr.Progress(track_tqdm=True)
 ) -> pd.DataFrame:
-    """
-    Main function to process images, run them through selected models,
-    and return results as a Pandas DataFrame.
-    """
+    """Main function to process images and return results as a Pandas DataFrame."""
     if not files:
         gr.Warning("No images uploaded. Please upload files to evaluate.")
         return pd.DataFrame()
-        
     if not selected_model_names:
         gr.Warning("No models selected. Please select at least one model.")
         return pd.DataFrame()
 
     try:
         image_paths = [Path(f.name) for f in files]
-        all_results = []
-        scorers = get_scorers(selected_model_names)
-        
-        # Use a single tqdm instance for progress tracking
-        pbar = tqdm(total=len(image_paths), desc="Processing images")
+        all_results, scorers = [], get_scorers(selected_model_names)
         
-        for i in range(0, len(image_paths), batch_size):
+        for i in tqdm(range(0, len(image_paths), batch_size), desc="Processing Batches"):
             batch_paths = image_paths[i : i + batch_size]
-            
-            # Load images for the current batch
             try:
                 batch_images = [Image.open(p).convert("RGB") for p in batch_paths]
             except Exception as e:
                 gr.Warning(f"Skipping a batch due to an error loading an image: {e}")
-                pbar.update(len(batch_paths))
                 continue
-
-            # Get scores from all selected models for the batch
-            batch_scores: Dict[str, List[float]] = {}
-            for scorer in scorers:
-                batch_scores[scorer.model_name] = scorer.score_batch(batch_images)
             
-            # Collate results for the batch
+            batch_scores = {scorer.model_name: scorer.score_batch(batch_images) for scorer in scorers}
+            
             for j, path in enumerate(batch_paths):
-                result_row = {"Image": Image.open(path), "Filename": path.name}
-                
-                scores_for_avg = []
+                result_row = {"Image": str(path), "Filename": path.name}
+                scores_for_avg = [batch_scores[s.model_name][j] for s in scorers]
                 for scorer in scorers:
-                    score = batch_scores[scorer.model_name][j]
-                    result_row[scorer.model_name] = score
-                    scores_for_avg.append(score)
-                
-                # Calculate average score
-                if scores_for_avg:
-                    result_row["Average Score"] = round(np.mean(scores_for_avg), 4)
-                else:
-                    result_row["Average Score"] = 0.0
-                
+                    result_row[scorer.model_name] = batch_scores[scorer.model_name][j]
+                result_row["Average Score"] = round(np.mean(scores_for_avg), 4) if scores_for_avg else 0.0
                 all_results.append(result_row)
-            
-            pbar.update(len(batch_paths))
 
-        pbar.close()
-        
-        if not all_results:
-            gr.Warning("Processing completed, but no results were generated.")
-            return pd.DataFrame()
-            
-        return pd.DataFrame(all_results)
+        return pd.DataFrame(all_results) if all_results else pd.DataFrame()
         
     except Exception as e:
         gr.Error(f"A critical error occurred: {e}")
-        # Clean up in case of failure
-        for scorer in _loaded_models_cache.values():
-            scorer.release_model()
-        _loaded_models_cache.clear()
         return pd.DataFrame()
 
-
 # ==================================================================================
 # 3. GRADIO USER INTERFACE
 # ==================================================================================
 
 def create_ui() -> gr.Blocks:
     """Creates and configures the Gradio web interface."""
-    
     all_model_names = list(MODEL_REGISTRY.keys())
-    
-    # Define headers and datatypes for the results table
     dataframe_headers = ["Image", "Filename"] + all_model_names + ["Average Score"]
     dataframe_datatypes = ["image", "str"] + ["number"] * (len(all_model_names) + 1)
     
-    with gr.Blocks(theme=gr.themes.Soft(), title="Image Aesthetic Scorer") as demo:
-        gr.Markdown(
-            """
-            # 🖼️ Modern Image Aesthetic Scorer
-            Upload your images, select the scoring models, and click 'Evaluate'.
-            The results table supports **interactive sorting** (click on headers) and can be **downloaded as a CSV**.
-            """
-        )
+    with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Image Aesthetic Scorer") as demo:
+        gr.Markdown("# 🖼️ Modern Image Aesthetic Scorer")
+        gr.Markdown("Upload images, select models, and click 'Evaluate'. Results table supports **interactive sorting** and **downloading as CSV**.")
 
         with gr.Row():
             with gr.Column(scale=1):
-                gr.Markdown("### ⚙️ Settings")
-                input_files = gr.Files(
-                    label="Upload Images",
-                    file_count="multiple",
-                    file_types=["image"],
-                )
-                
-                with gr.Accordion("Advanced Configuration", open=False):
-                    model_checkboxes = gr.CheckboxGroup(
-                        choices=all_model_names,
-                        value=all_model_names,
-                        label="Scoring Models",
-                        info="Choose which models to use for evaluation.",
-                    )
-                    batch_size_slider = gr.Slider(
-                        minimum=1,
-                        maximum=64,
-                        value=8,
-                        step=1,
-                        label="Batch Size",
-                        info="Adjust based on your VRAM. Higher is faster.",
-                    )
-
+                input_files = gr.Files(label="Upload Images", file_count="multiple", file_types=["image"])
+                model_checkboxes = gr.CheckboxGroup(choices=all_model_names, value=all_model_names, label="Scoring Models")
+                batch_size_slider = gr.Slider(minimum=1, maximum=64, value=8, step=1, label="Batch Size", info="Adjust based on your VRAM.")
                 with gr.Row():
                     process_button = gr.Button("🚀 Evaluate Images", variant="primary")
                     clear_button = gr.Button("🧹 Clear All")
             
             with gr.Column(scale=3):
-                gr.Markdown("### 📊 Results")
+                # CORRECTED LINE: height and show_download_button are passed directly here.
                 results_dataframe = gr.DataFrame(
                     headers=dataframe_headers,
                     datatype=dataframe_datatypes,
                     label="Evaluation Scores",
                     interactive=True,
-                    # Enable the download button directly on the component
+                    height=800,
+                    show_download_button=True
                 )
-                # This is a cleaner way to show the download button
-                results_dataframe.style(height=800, show_download_button=True)
-
 
-        # --- Event Handlers ---
         process_button.click(
             fn=evaluate_images,
             inputs=[input_files, model_checkboxes, batch_size_slider],
-            outputs=[results_dataframe],
-            concurrency_limit=1 # Only one evaluation at a time
+            outputs=[results_dataframe]
         )
 
         def clear_outputs():
-            # Release all models from memory when clearing
-            for scorer in _loaded_models_cache.values():
+            for scorer in list(_loaded_models_cache.values()):
                 scorer.release_model()
             _loaded_models_cache.clear()
             gr.Info("Cleared results and released models from memory.")
-            # Return an empty DataFrame to clear the table
-            return pd.DataFrame()
-
-        clear_button.click(
-            fn=clear_outputs,
-            inputs=[],
-            outputs=[results_dataframe],
-        )
+            return pd.DataFrame(), None # Clear dataframe and file input
         
+        clear_button.click(fn=clear_outputs, outputs=[results_dataframe, input_files])
     return demo
 
 # ==================================================================================
@@ -400,8 +314,6 @@ def create_ui() -> gr.Blocks:
 # ==================================================================================
 
 if __name__ == "__main__":
-    # Ensure cache directory exists
     os.makedirs(CACHE_DIR, exist_ok=True)
-    
     app = create_ui()
-    app.queue().launch(share=False)
+    app.queue().launch()