Update app.py

app.py

@@ -1,6 +1,9 @@
 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
@@ -9,15 +12,31 @@ from io import BytesIO
 # Cấu hình API Key
 api_key = os.getenv("GEMINI_API_KEY")
 if not api_key:
-
     raise ValueError("⚠️ GEMINI_API_KEY is missing!")
 
 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()
+
+def upscale_image(image):
+    """Nâng cấp độ phân giải ảnh bằng SRCNN"""
+    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 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_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
@@ -39,6 +58,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
             images.append(img)
     return images
 
@@ -65,8 +85,8 @@ 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",
-    description="Upload ảnh và nhập yêu cầu chỉnh sửa. Chọn số lượng ảnh muốn tạo (1-4).",
+    title="Chỉnh sửa ảnh bằng Gemini AI + SRCNN",
+    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ý.",
 )
 
 demo.launch()