Update app.py

app.py

@@ -31,22 +31,16 @@ def load_image_as_bytes(image_path):
     except UnidentifiedImageError:
         print(f"❌ Lỗi: Không thể mở file {image_path} (định dạng không hợp lệ)")
         return None
-
     except Exception as e:
         print(f"❌ Lỗi khi đọc ảnh {image_path}: {e}")
         return None
 
-async def generate_image(image_bytes_list, text_input, max_retries=5):
-    """Gửi request và nhận kết quả từ Gemini API (Xử lý lỗi 429 với popup thông báo)"""
-    attempt = 0  # Đếm số lần thử lại
-    while attempt < max_retries:
+async def generate_image(image_bytes_list, text_input):
+    """Gửi request và nhận kết quả từ Gemini API (Xử lý lỗi 429)"""
+    while True:
         try:
             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 len(image_parts) == 1 else
-                sum([[f"Image {idx+1}:", part] for idx, part in enumerate(image_parts)], []) if len(image_parts) == 2 else []
-            )
+            contents = [text_input, image_parts] if image_parts else [text_input]
 
             response = await asyncio.to_thread(
                 client.models.generate_content,
@@ -56,9 +50,8 @@ async def generate_image(image_bytes_list, text_input, max_retries=5):
             )
 
             if not response or not response.candidates or not response.candidates[0].content:
-                attempt += 1
-                gr.Warning(f"⚠️ API trả về dữ liệu không hợp lệ (thử {attempt}/{max_retries}). Đang thử lại...")
-                continue  # Thử lại ngay lập tức
+                print("❌ Lỗi: Phản hồi API không hợp lệ")
+                return []
 
             images = []
             for part in response.candidates[0].content.parts:
@@ -67,37 +60,33 @@ async def generate_image(image_bytes_list, text_input, max_retries=5):
                         img = Image.open(BytesIO(part.inline_data.data))
                         images.append(img)
                     except Exception as e:
-                        gr.Error(f"❌ Lỗi khi hiển thị ảnh: {e}")
-            return images  # Thành công, trả về ảnh ngay lập tức
+                        print(f"❌ Lỗi khi hiển thị ảnh: {e}")
+            return images
 
         except Exception as e:
             error_message = str(e)
-            if "RESOURCE_EXHAUSTED" in error_message:
+            if "429" in error_message and "RESOURCE_EXHAUSTED" in error_message:
                 try:
                     # Trích xuất retryDelay từ JSON lỗi
                     error_json = json.loads(error_message.split("RESOURCE_EXHAUSTED. ")[1])
                     retry_delay = int(error_json["error"]["details"][-1]["retryDelay"][:-1])  # Lấy số giây từ '2s'
 
-                    gr.Warning(f"⚠️ API quá tải! Sẽ thử lại sau {retry_delay} giây...")
-                    
-                    # Đếm ngược thời gian chờ (có popup)
+                    print(f"⚠️ Đã vượt quá hạn mức! Chờ {retry_delay} giây trước khi thử lại...")
+
+                    # Đếm ngược
                     for i in range(retry_delay, 0, -1):
-                        gr.Info(f"⏳ Thử lại sau {i} giây...")
+                        print(f"⏳ Thử lại sau {i} giây...", end="\r")
                         time.sleep(1)
 
-                    gr.Info("🔄 Đang thử lại...")
-                    attempt += 1
-                    continue  # Thử lại sau thời gian chờ
+                    print("\n🔄 Đang thử lại...")
+                    continue  # Thử lại request
 
                 except Exception as parse_error:
-                    gr.Error(f"❌ Lỗi khi phân tích retryDelay: {parse_error}")
+                    print(f"❌ Lỗi khi phân tích retryDelay: {parse_error}")
 
-            gr.Error(f"❌ Lỗi khi gọi API Gemini: {e}")
+            print(f"❌ Lỗi khi gọi API Gemini: {e}")
             return []
 
-    gr.Error("❌ Đã thử lại nhiều lần nhưng vẫn thất bại!")
-    return []
-
 async def process_request(images, text, num_requests):
     """Chạy nhiều request song song"""
     image_bytes_list = [load_image_as_bytes(image) if image else None for image in images]
@@ -110,14 +99,13 @@ async def process_request(images, text, num_requests):
     generated_images = [img for result in results for img in result]
     
     # Resize ảnh giữ nguyên tỷ lệ chiều cao
-    resized_images = [img.resize((3840, int(img.height * (3840 / img.width))), Image.LANCZOS) for img in generated_images]
+    resized_images = [img.resize((2560, int(img.height * (2560 / img.width))), Image.LANCZOS) for img in generated_images]
     print("num_requests", num_requests)
     print("tasks", len(tasks))
     print("generated_images", len(generated_images))
     print("resized_images", len(resized_images))
     
-    # return generated_images + resized_images
-    return resized_images
+    return generated_images + resized_images
 
 def gradio_interface(image1, image2, text, num_requests):
     """Hàm Gradio xử lý yêu cầu và trả về ảnh"""
@@ -134,8 +122,8 @@ demo = gr.Interface(
         gr.Slider(minimum=1, maximum=8, step=1, value=4, label="Số lượng ảnh cần tạo")  # Tăng lên 8
     ],
     outputs=gr.Gallery(label="Kết quả chỉnh sửa", columns=4),
-    title="Chỉnh sửa ảnh bằng Gemini AI + Upscale",
-    description="Upload tối đa 2 ảnh và nhập yêu cầu chỉnh sửa. Hiển thị ảnh gốc từ API và đã qua upscale.",
+    title="Chỉnh sửa ảnh bằng Gemini AI + SRCNN",
+    description="Upload tối đa 2 ảnh và nhập yêu cầu chỉnh sửa. Hiển thị ảnh gốc từ API và ảnh đã qua SRCNN.",
 )
 
 demo.launch()