Update app.py

app.py

@@ -1,4 +1,56 @@
 import gradio as gr
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import ToTensor
+import torch.nn.functional as F
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+        # Definimos las capas convolucionales
+        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
+        self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
+
+        # Definimos capas fully connected
+        self.fc1 = nn.Linear(128 * 3 * 3, 256)
+        self.fc2 = nn.Linear(256, 10)
+
+        # Definimos un max pooling y dropout
+        self.pool = nn.MaxPool2d(2, 2)
+        self.dropout = nn.Dropout(0.25)
+
+    def forward(self, x):
+        # Pasamos las entradas por las capas convolucionales y el max pooling
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = self.pool(F.relu(self.conv3(x)))
+        
+        # Aplanamos la salida de las capas convolucionales para pasar a fully connected
+        x = x.view(-1, 128 * 3 * 3)
+        
+        # Pasamos por las capas fully connected
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = self.fc2(x)
+        
+        return x
+
+model = CNN().to(device)
+
+# Cargar el modelo en la CPU
+model = CNN().to(device)
+model.load_state_dict(torch.load("model_mnist_cnn.pth", map_location=torch.device('cpu')))
 
 def predict(im):