Update app.py

app.py

@@ -2,15 +2,12 @@ import gradio as gr
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
-from transformers import ViTForImageClassification, AutoTokenizer
+
 
 
 # Nombre del modelo en el repositorio de Hugging Face
 model_name = "mkjaramillo/cancer"
-
-# Cargar el tokenizer y el modelo desde el repositorio de Hugging Face
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = ViTForImageClassification.from_pretrained(model_name)
+model = torch.load("model_name")
 
 # Transformación de la imagen
 image_transform = transforms.Compose([
@@ -31,10 +28,10 @@ def classify_image(image):
     outputs = model(image)
     
     # Obtener las predicciones
-    predictions = torch.argmax(outputs.logits, dim=1)
+    predictions = torch.argmax(outputs, dim=1)
     
     # Obtener la etiqueta de la predicción
-    label = tokenizer.decode(predictions.item())
+    label = predictions.item()
 
     # Retornar la etiqueta de la predicción
     return label