Update app.py

app.py

@@ -1,13 +1,11 @@
 import gradio as gr
 import asyncio
 import os
-import torch
-import torchvision.transforms as transforms
-from torchvision.utils import save_image
 from google import genai
 from google.genai import types
 from PIL import Image
 from io import BytesIO
+from super_image import RcanModel, ImageLoader
 
 # Cấu hình API Key
 api_key = os.getenv("GEMINI_API_KEY")
@@ -16,28 +14,20 @@ if not api_key:
 
 client = genai.Client(api_key=api_key)
 
-# Load SRCNN từ Torch Hub
-model = torch.hub.load('pytorch/vision:v0.10.0', 'srcnn', pretrained=True)
-model.eval()
+# Load RCAN-BAM model
+model = RcanModel.from_pretrained('eugenesiow/rcan-bam', scale=2)
 
-def upscale_image(image, target_resolution=(2560, 1440)):
-    """Nâng cấp độ phân giải ảnh bằng nội suy trước khi qua SRCNN"""
-    image = image.resize(target_resolution, Image.BICUBIC)  # Nội suy trước khi SRCNN xử lý
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Lambda(lambda x: x.unsqueeze(0))  # Thêm batch dimension
-    ])
-    img_tensor = transform(image)
-    with torch.no_grad():
-        upscaled_tensor = model(img_tensor)
-    upscaled_image = transforms.ToPILImage()(upscaled_tensor.squeeze(0))
-    return upscaled_image
+def upscale_image(image):
+    """Nâng cấp độ phân giải ảnh bằng RCAN-BAM"""
+    inputs = ImageLoader.load_image(image)
+    preds = model(inputs)
+    return ImageLoader.to_pil_image(preds)
 
 def load_image_as_bytes(image_path):
     """Chuyển ảnh thành dữ liệu nhị phân"""
     with Image.open(image_path) as img:
         img = img.convert("RGB")  # Đảm bảo ảnh là RGB
-        img = upscale_image(img)  # SRCNN trước khi gửi đi
+        img = upscale_image(img)  # RCAN-BAM xử lý
         img_bytes = BytesIO()
         img.save(img_bytes, format="JPEG")  # Lưu ảnh vào buffer
         return img_bytes.getvalue()  # Lấy dữ liệu nhị phân
@@ -59,7 +49,7 @@ async def generate_image(image_bytes, text_input):
     for part in response.candidates[0].content.parts:
         if part.inline_data is not None:
             img = Image.open(BytesIO(part.inline_data.data))
-            img = upscale_image(img)  # SRCNN sau khi nhận ảnh từ Gemini
+            img = upscale_image(img)  # RCAN-BAM sau khi nhận ảnh từ Gemini
             images.append(img)
     return images
 
@@ -86,7 +76,7 @@ demo = gr.Interface(
         gr.Slider(minimum=1, maximum=4, step=1, value=4, label="Số lượng ảnh cần tạo")
     ],
     outputs=gr.Gallery(label="Kết quả chỉnh sửa", columns=4),
-    title="Chỉnh sửa ảnh bằng Gemini AI + SRCNN",
+    title="Chỉnh sửa ảnh bằng Gemini AI + RCAN-BAM",
     description="Upload ảnh và nhập yêu cầu chỉnh sửa. Ảnh được nâng cấp độ phân giải trước và sau khi xử lý.",
 )