Update app.py

app.py

@@ -8,6 +8,7 @@ from transformers import ViTForImageClassification
 from torch import nn
 from torch.cuda.amp import autocast
 import os
+from contextlib import nullcontext
 
 # Global configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -19,7 +20,7 @@ label_mapping = {
     1: "Меланоцитарный невус",
     2: "Базальноклеточная карцинома",
     3: "Актинический кератоз",
-    4: "Доброкачественная кератоза",
+    4: "Доброкачественная кератома",
     5: "Дерматофиброма",
     6: "Сосудистые поражения"
 }
@@ -88,7 +89,9 @@ class ModelHandler:
             return {"error": "Модели не загружены"}
 
         inputs = transform_image(image)
-        with autocast():
+        # Handle autocast based on device
+        ctx = autocast() if device.type == 'cuda' else nullcontext()
+        with ctx:
             outputs = self.efficientnet(inputs)
             probs = torch.nn.functional.softmax(outputs, dim=1)
 
@@ -100,8 +103,9 @@ class ModelHandler:
             return {"error": "Модели не загружены"}
 
         inputs = transform_image(image)
-        with autocast():
-            outputs = self.deit(inputs).logits
+        ctx = autocast() if device.type == 'cuda' else nullcontext()
+        with ctx:
+            outputs = self.deit(pixel_values=inputs).logits  # Corrected parameter
             probs = torch.nn.functional.softmax(outputs, dim=1)
 
         return self._format_predictions(probs)
@@ -112,23 +116,22 @@ class ModelHandler:
             return {"error": "Модели не загружены"}
 
         inputs = transform_image(image)
-        with autocast():
+        ctx = autocast() if device.type == 'cuda' else nullcontext()
+        with ctx:
             eff_probs = torch.nn.functional.softmax(self.efficientnet(inputs), dim=1)
-            deit_probs = torch.nn.functional.softmax(self.deit(inputs).logits, dim=1)
+            deit_probs = torch.nn.functional.softmax(self.deit(pixel_values=inputs).logits, dim=1)
             ensemble_probs = (eff_probs + deit_probs) / 2
 
         return self._format_predictions(ensemble_probs)
 
-        def _format_predictions(self, probs):
-            top5_probs, top5_indices = torch.topk(probs, 5)
-            result = {}
-            for i in range(5):
-                idx = top5_indices[0][i].item()
-                label = label_mapping.get(idx, f"Класс {idx}")
-                # return raw prob, not percent:
-                result[label] = float(top5_probs[0][i].item())
-            return result
-
+    def _format_predictions(self, probs):  # Corrected indentation
+        top5_probs, top5_indices = torch.topk(probs, 5)
+        result = {}
+        for i in range(5):
+            idx = top5_indices[0][i].item()
+            label = label_mapping.get(idx, f"Класс {idx}")
+            result[label] = float(top5_probs[0][i].item())
+        return result
 
 # Initialize model handler
 model_handler = ModelHandler()
@@ -184,4 +187,4 @@ def create_interface():
 if __name__ == "__main__":
     interface = create_interface()
     print("🚀 Запуск интерфейса...")
-    interface.launch(share=True)
+    interface.launch(server_port=7860)  # Explicitly set port if needed