Update app.py

app.py

@@ -3,12 +3,12 @@ import asyncio
 import os
 import torch
 import torch.nn as nn
-import torch.optim as optim
 import torchvision.transforms as transforms
 from google import genai
 from google.genai import types
 from PIL import Image
 from io import BytesIO
+from huggingface_hub import hf_hub_download
 
 # Cấu hình API Key
 api_key = os.getenv("GEMINI_API_KEY")
@@ -18,23 +18,19 @@ if not api_key:
 client = genai.Client(api_key=api_key)
 
 # Định nghĩa mô hình SRCNN
-class SRCNN(nn.Module):
-    def __init__(self):
-        super(SRCNN, self).__init__()
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)
-        self.conv2 = nn.Conv2d(64, 32, kernel_size=5, padding=2)
-        self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2)
-        self.relu = nn.ReLU()
-    
-    def forward(self, x):
-        x = self.relu(self.conv1(x))
-        x = self.relu(self.conv2(x))
-        x = self.conv3(x)
-        return x
+def SRCNN():
+    return nn.Sequential(
+        nn.Conv2d(3, 64, kernel_size=9, padding=4),
+        nn.ReLU(inplace=True),
+        nn.Conv2d(64, 32, kernel_size=5, padding=2),
+        nn.ReLU(inplace=True),
+        nn.Conv2d(32, 3, kernel_size=5, padding=2)
+    )
 
-# Khởi tạo mô hình SRCNN
+# Khởi tạo mô hình SRCNN và tải trọng số từ Hugging Face
 model = SRCNN()
-model.load_state_dict(torch.load("srcnn.pth", map_location=torch.device('cpu')))
+pth_path = hf_hub_download(repo_id="sunbv56/srcnn", filename="srcnn.pth")
+model.load_state_dict(torch.load(pth_path, map_location=torch.device('cpu')))
 model.eval()
 
 def upscale_image(image):