Update app.py

app.py

@@ -47,32 +47,74 @@ client = genai.Client(api_key=api_key)
 #     return output_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_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
+    """Chuyển ảnh thành dữ liệu nhị phân với kiểm tra lỗi"""
+    try:
+        with Image.open(image_path) as img:
+            img = img.convert("RGB")  # Đảm bảo ảnh là RGB
+            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
+    except FileNotFoundError:
+        print(f"❌ Lỗi: Không tìm thấy file {image_path}")
+        return None
+    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):
-    """Gửi request và nhận kết quả từ Gemini API"""
-    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,
-        model="gemini-2.0-flash-exp-image-generation",
-        contents=contents,
-        config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
-    )
-    
-    images = []
-    if response and response.candidates and response.candidates[0].content:
-        for part in response.candidates[0].content.parts:
-            if part.inline_data is not None:
-                img = Image.open(BytesIO(part.inline_data.data))
-                images.append(img)
-    return images
+    """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 image_parts else [text_input]
+
+            response = await asyncio.to_thread(
+                client.models.generate_content,
+                model="gemini-2.0-flash-exp-image-generation",
+                contents=contents,
+                config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
+            )
+
+            if not response or not response.candidates or not response.candidates[0].content:
+                print("❌ Lỗi: Phản hồi API không hợp lệ")
+                return []
+
+            images = []
+            for part in response.candidates[0].content.parts:
+                if part.inline_data is not None:
+                    try:
+                        img = Image.open(BytesIO(part.inline_data.data))
+                        images.append(img)
+                    except Exception as e:
+                        print(f"❌ Lỗi khi hiển thị ảnh: {e}")
+            return images
+
+        except Exception as e:
+            error_message = str(e)
+            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'
+
+                    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):
+                        print(f"⏳ Thử lại sau {i} giây...", end="\r")
+                        time.sleep(1)
+
+                    print("\n🔄 Đang thử lại...")
+                    continue  # Thử lại request
+
+                except Exception as parse_error:
+                    print(f"❌ Lỗi khi phân tích retryDelay: {parse_error}")
+
+            print(f"❌ Lỗi khi gọi API Gemini: {e}")
+            return []
 
 async def process_request(images, text, num_requests):
     """Chạy nhiều request song song"""
@@ -87,10 +129,11 @@ async def process_request(images, text, num_requests):
     
     # Resize ảnh giữ nguyên tỷ lệ chiều cao
     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))
+    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
 
 def gradio_interface(image1, image2, text, num_requests):
@@ -98,14 +141,14 @@ def gradio_interface(image1, image2, text, num_requests):
     images = [img for img in [image1, image2] if img]
     return asyncio.run(process_request(images, text, num_requests))
 
-# Tạo giao diện Gradio
+# Tạo giao diện Gradio với slider từ 1 đến 8
 demo = gr.Interface(
     fn=gradio_interface,
     inputs=[
         gr.Image(type='filepath', label="Upload hình ảnh 1"),
         gr.Image(type='filepath', label="Upload hình ảnh 2"),
         gr.Textbox(label="Nhập yêu cầu chỉnh sửa hình ảnh"),
-        gr.Slider(minimum=1, maximum=4, step=1, value=4, label="Số lượng ảnh cần tạo")
+        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 + SRCNN",