Update app.py

app.py

@@ -17,33 +17,34 @@ if not api_key:
 
 client = genai.Client(api_key=api_key)
 
-# Định nghĩa mô hình SRCNN
-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)
-    )
+# # Định nghĩa mô hình SRCNN
+# 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 và tải trọng số từ Hugging Face
-model = SRCNN()
-pth_path = hf_hub_download(repo_id="sunbv56/srcnn", filename="srcnn_model.pth")
-model.load_state_dict(torch.load(pth_path, map_location=torch.device('cpu')))
-model.eval()
+# # Khởi tạo mô hình SRCNN và tải trọng số từ Hugging Face
+# model = SRCNN()
+# pth_path = hf_hub_download(repo_id="sunbv56/srcnn", filename="srcnn_model.pth")
+# model.load_state_dict(torch.load(pth_path, map_location=torch.device('cpu')))
+# 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
-    ])
-    input_tensor = transform(image)
-    with torch.no_grad():
-        output_tensor = model(input_tensor)
-    output_image = transforms.ToPILImage()(output_tensor.squeeze(0))
-    return output_image
+# 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
+#     ])
+#     input_tensor = transform(image)
+#     with torch.no_grad():
+#         output_tensor = model(input_tensor)
+#     output_image = transforms.ToPILImage()(output_tensor.squeeze(0))
+#     return output_image
 
 def load_image_as_bytes(image_path):
     """Chuyển ảnh thành dữ liệu nhị phân"""
@@ -53,11 +54,10 @@ def load_image_as_bytes(image_path):
         img.save(img_bytes, format="JPEG")  # Lưu ảnh vào buffer
         return img_bytes.getvalue()  # Lấy dữ liệu nhị phân
 
-async def generate_image(image_bytes, text_input):
+async def generate_image(image_bytes_list, text_input):
     """Gửi request và nhận kết quả từ Gemini API"""
-    contents = [text_input]
-    if image_bytes:
-        contents.append(types.Part(inline_data=types.Blob(data=image_bytes, mime_type="image/jpeg")))
+    image_parts = [types.Part(inline_data=types.Blob(data=img, mime_type="image/jpeg")) for img in image_bytes_list if img]
+    contents = [text_input, image_parts] if image_parts else [text_input]
     
     response = await asyncio.to_thread(
         client.models.generate_content,
@@ -72,9 +72,6 @@ async def generate_image(image_bytes, text_input):
             if part.inline_data is not None:
                 img = Image.open(BytesIO(part.inline_data.data))
                 images.append(img)
-    else:
-        print("⚠️ Gemini API không trả về kết quả hợp lệ!")
-    
     return images
 
 async def process_request(image, text, num_requests):
@@ -85,10 +82,11 @@ async def process_request(image, text, num_requests):
     
     # Hợp nhất danh sách ảnh từ các request
     original_images = [img for result in results for img in result]
-    resized_images = [img.resize((2560, 1440), Image.LANCZOS) for img in original_images]  # Resize trước khi upscale
-    srcnn_images = [upscale_image(img) for img in resized_images]
+    resized_images = [img.resize((2560, int(img.height * (2560 / img.width)))), Image.LANCZOS) for img in original_images]  # Resize trước khi upscale
+    # srcnn_images = [upscale_image(img) for img in resized_images]
     
-    return original_images + srcnn_images  # 4 ảnh gốc + 4 ảnh đã qua SRCNN
+    return resized_images  # 4 ảnh gốc
+    # return resized_images + srcnn_images  # 4 ảnh gốc + 4 ảnh đã qua SRCNN
 
 def gradio_interface(image, text, num_requests):
     """Hàm Gradio xử lý yêu cầu và trả về ảnh"""