Update app.py

app.py

@@ -1,28 +1,33 @@
-import gradio as gr
+# app.py
+import os
+import time
+import tempfile
+import threading
+import queue
+import logging
+import numpy as np  # 追加: 画像処理の最適化用
+from io import BytesIO
+from PIL import Image
+from pydantic import BaseModel
 from fastapi import FastAPI, HTTPException, Body
 from fastapi.responses import StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
+import gradio as gr
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
-from PIL import Image
-from io import BytesIO
-import tempfile
-import time
-import os
-import logging
-import numpy as np  # 追加: 画像処理の最適化用
-import threading  # 追加: 並列処理のため
-import queue  # 追加: WebDriverプール用
 from concurrent.futures import ThreadPoolExecutor  # 追加: 並列処理用
 from huggingface_hub import hf_hub_download
 
-# 正しいGemini関連のインポート
-import google.generativeai as genai
+# 既存の Gemini ライブラリ
+import google.generativeai as genai_old
+
+# 新しい Gemini ライブラリ（2.5系モデル用）
+from google import genai as genai_new
+from google.genai import types
 
 # ロギング設定
 logging.basicConfig(level=logging.INFO)
@@ -56,7 +61,6 @@ class WebDriverPool:
                 options.add_argument("--disable-features=NetworkService")
                 options.add_argument("--dns-prefetch-disable")
                 
-                # 環境変数からWebDriverパスを取得（任意）
                 webdriver_path = os.environ.get("CHROMEDRIVER_PATH")
                 if webdriver_path and os.path.exists(webdriver_path):
                     logger.info(f"CHROMEDRIVER_PATH使用: {webdriver_path}")
@@ -66,7 +70,6 @@ class WebDriverPool:
                     logger.info("デフォルトのChromeDriverを使用")
                     return webdriver.Chrome(options=options)
         
-        # 最大数に達した場合は待機
         logger.info("WebDriverプールがいっぱいです。利用可能なドライバーを待機中...")
         return self.driver_queue.get()
     
@@ -74,7 +77,6 @@ class WebDriverPool:
         """ドライバーをプールに戻す"""
         if driver:
             try:
-                # ブラウザをリセット
                 driver.get("about:blank")
                 driver.execute_script("""
                     document.documentElement.style.overflow = '';
@@ -107,90 +109,70 @@ class WebDriverPool:
             self.count = 0
 
 # グローバルなWebDriverプールを作成
-# サーバー環境のリソースに合わせて調整
 driver_pool = WebDriverPool(max_drivers=int(os.environ.get("MAX_WEBDRIVERS", "3")))
 
-# --- Gemini統合 ---
+# --- リクエストモデル ---
 class GeminiRequest(BaseModel):
     """Geminiへのリクエストデータモデル"""
     text: str
     extension_percentage: float = 10.0  # デフォルト値10%
-    temperature: float = 0.5  # デフォルト値を0.5に設定
-    trim_whitespace: bool = True  # 余白トリミングオプション（デフォルト有効）
-    style: str = "standard"  # デフォルトはstandard
+    temperature: float = 0.5            # デフォルト値を0.5
+    trim_whitespace: bool = True        # 余白トリミングオプション（既定で有効）
+    style: str = "standard"             # デフォルトはstandard
 
 class ScreenshotRequest(BaseModel):
     """スクリーンショットリクエストモデル"""
     html_code: str
-    extension_percentage: float = 10.0  # デフォルト値10%
-    trim_whitespace: bool = True  # 余白トリミングオプション（デフォルト有効）
-    style: str = "standard"  # デフォルトはstandard
+    extension_percentage: float = 10.0
+    trim_whitespace: bool = True
+    style: str = "standard"
 
-# HTMLのFont Awesomeレイアウトを改善する関数 - プリロード機能を追加
+# --- Font Awesomeレイアウト改善 ---
 def enhance_font_awesome_layout(html_code):
     """Font Awesomeレイアウトを改善し、プリロードタグを追加"""
-    # Font Awesomeリソースのプリロード - パフォーマンス向上
     fa_preload = """
     <link rel="preload" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/webfonts/fa-solid-900.woff2" as="font" type="font/woff2" crossorigin>
     <link rel="preload" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/webfonts/fa-regular-400.woff2" as="font" type="font/woff2" crossorigin>
     <link rel="preload" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/webfonts/fa-brands-400.woff2" as="font" type="font/woff2" crossorigin>
     """
-    
-    # CSSを追加
     fa_fix_css = """
     <style>
-    /* Font Awesomeアイコンのレイアウト���正 */
     [class*="fa-"] {
         display: inline-block !important;
         margin-right: 8px !important;
         vertical-align: middle !important;
     }
-    
-    /* テキスト内のアイコン位置調整 */
-    h1 [class*="fa-"], h2 [class*="fa-"], h3 [class*="fa-"], 
+    h1 [class*="fa-"], h2 [class*="fa-"], h3 [class*="fa-"],
     h4 [class*="fa-"], h5 [class*="fa-"], h6 [class*="fa-"] {
         vertical-align: middle !important;
         margin-right: 10px !important;
     }
-    
-    /* 特定パターンの修正 */
     .fa + span, .fas + span, .far + span, .fab + span,
     span + .fa, span + .fas, span + .far + span {
         display: inline-block !important;
         margin-left: 5px !important;
     }
-    
-    /* カード内アイコン修正 */
     .card [class*="fa-"], .card-body [class*="fa-"] {
         float: none !important;
         clear: none !important;
         position: relative !important;
     }
-    
-    /* アイコンと文字が重なる場合の調整 */
     li [class*="fa-"], p [class*="fa-"] {
         margin-right: 10px !important;
     }
-    
-    /* インラインアイコンのスペーシング */
     .inline-icon {
         display: inline-flex !important;
         align-items: center !important;
         justify-content: flex-start !important;
     }
-    
-    /* アイコン後のテキスト */
     [class*="fa-"] + span {
         display: inline-block !important;
         vertical-align: middle !important;
     }
     </style>
     """
-
-    # headタグがある場合はその中に追加
     if '<head>' in html_code:
         return html_code.replace('</head>', f'{fa_preload}{fa_fix_css}</head>')
-    # HTMLタグがある場合はその後に追加
     elif '<html' in html_code:
         head_end = html_code.find('</head>')
         if head_end > 0:
@@ -199,282 +181,175 @@ def enhance_font_awesome_layout(html_code):
             body_start = html_code.find('<body')
             if body_start > 0:
                 return html_code[:body_start] + f'<head>{fa_preload}{fa_fix_css}</head>' + html_code[body_start:]
-
-    # どちらもない場合は先頭に追加
     return f'<html><head>{fa_preload}{fa_fix_css}</head>' + html_code + '</html>'
 
+# --- システムインストラクション読み込み ---
 def load_system_instruction(style="standard"):
     """
     指定されたスタイルのシステムインストラクションを読み込む
-
-    Args:
-        style: 使用するスタイル名 (standard, cute, resort, cool, dental, school)
-
-    Returns:
-        読み込まれたシステムインストラクション
     """
+    valid_styles = ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"]
+    if style not in valid_styles:
+        logger.warning(f"無効なスタイル '{style}' が指定されました。デフォルトの 'standard' を使用します。")
+        style = "standard"
+    logger.info(f"スタイル '{style}' のシステムインストラクションを読み込みます")
+
+    # ローカルファイル優先
+    local_path = os.path.join(os.path.dirname(__file__), style, "prompt.txt")
+    if os.path.exists(local_path):
+        logger.info(f"ローカルファイルを使用: {local_path}")
+        with open(local_path, 'r', encoding='utf-8') as file:
+            return file.read()
+
+    # HuggingFace から取得
     try:
-        # 有効なスタイル一覧
-        valid_styles = ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"]
-
-        # スタイルの検証
-        if style not in valid_styles:
-            logger.warning(f"無効なスタイル '{style}' が指定されました。デフォルトの 'standard' を使用します。")
-            style = "standard"
-
-        logger.info(f"スタイル '{style}' のシステムインストラクションを読み込みます")
-
-        # まず、ローカルのスタイルディレクトリ内のprompt.txtを確認
-        local_path = os.path.join(os.path.dirname(__file__), style, "prompt.txt")
-
-        # ローカルファイルが存在する場合はそれを使用
-        if os.path.exists(local_path):
-            logger.info(f"ローカルファイルを使用: {local_path}")
-            with open(local_path, 'r', encoding='utf-8') as file:
-                instruction = file.read()
-            return instruction
-
-        # HuggingFaceリポジトリからのファイル読み込みを試行
-        try:
-            # スタイル固有のファイルパスを指定
-            file_path = hf_hub_download(
-                repo_id="tomo2chin2/GURAREKOstlyle",
-                filename=f"{style}/prompt.txt",
-                repo_type="dataset"
-            )
-
-            logger.info(f"スタイル '{style}' のプロンプトをHuggingFaceから読み込みました: {file_path}")
-            with open(file_path, 'r', encoding='utf-8') as file:
-                instruction = file.read()
-            return instruction
-
-        except Exception as style_error:
-            # スタイル固有ファイルの読み込みに失敗した場合、デフォルトのprompt.txtを使用
-            logger.warning(f"スタイル '{style}' のプロンプト読み込みに失敗: {str(style_error)}")
-            logger.info("デフォルトのprompt.txtを読み込みます")
-
-            file_path = hf_hub_download(
-                repo_id="tomo2chin2/GURAREKOstlyle",
-                filename="prompt.txt",
-                repo_type="dataset"
-            )
-
-            with open(file_path, 'r', encoding='utf-8') as file:
-                instruction = file.read()
-
-            logger.info("デフォルトのシステムインストラクションを読み込みました")
-            return instruction
-
-    except Exception as e:
-        error_msg = f"システムインストラクションの読み込みに失敗: {str(e)}"
-        logger.error(error_msg)
-        raise ValueError(error_msg)
+        file_path = hf_hub_download(
+            repo_id="tomo2chin2/GURAREKOstlyle",
+            filename=f"{style}/prompt.txt",
+            repo_type="dataset"
+        )
+        logger.info(f"HuggingFace から読み込み: {file_path}")
+        with open(file_path, 'r', encoding='utf-8') as file:
+            return file.read()
+    except Exception as style_error:
+        logger.warning(f"スタイル '{style}' の読み込み失敗: {style_error}")
+        logger.info("デフォルトの prompt.txt を読み込みます")
+        file_path = hf_hub_download(
+            repo_id="tomo2chin2/GURAREKOstlyle",
+            filename="prompt.txt",
+            repo_type="dataset"
+        )
+        with open(file_path, 'r', encoding='utf-8') as file:
+            return file.read()
 
+# --- テキストからHTML生成 ---
 def generate_html_from_text(text, temperature=0.5, style="standard"):
-    """テキストからHTMLを生成する"""
+    """
+    テキストからHTMLを生成する
+    gemini-2.5-flash-preview-04-17 のときのみ新ライブラリ＋thinkingBudget=0
+    """
     try:
-        # APIキーの取得と設定
         api_key = os.environ.get("GEMINI_API_KEY")
         if not api_key:
             logger.error("GEMINI_API_KEY 環境変数が設定されていません")
-            raise ValueError("GEMINI_API_KEY 環境変数が設定されていません")
+            raise ValueError("GEMINI_API_KEY が設定されていません")
 
-        # モデル名の取得（環境変数から、なければデフォルト値）
         model_name = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro")
-        logger.info(f"使用するGeminiモデル: {model_name}")
-
-        # Gemini APIの設定
-        genai.configure(api_key=api_key)
-
-        # 指定されたスタイルのシステムインストラクションを読み込む
-        system_instruction = load_system_instruction(style)
-
-        # モデル初期化
-        logger.info(f"Gemini APIにリクエストを送信: テキスト長さ = {len(text)}, 温度 = {temperature}, スタイル = {style}")
-
-        # モデル初期化
-        model = genai.GenerativeModel(model_name)
-
-        # 生成設定 - ばらつきを減らすために設定を調整
-        generation_config = {
-            "temperature": temperature,  # より低い温度を設定
-            "top_p": 0.7,                # 0.95から0.7に下げて出力の多様性を制限
-            "top_k": 20,                 # 64から20に下げて候補を絞る
-            "max_output_tokens": 8192,
-            "candidate_count": 1         # 候補は1つだけ生成
-        }
-
-        # 安全設定 - デフォルトの安全設定を使用
-        safety_settings = [
-            {
-                "category": "HARM_CATEGORY_HARASSMENT",
-                "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-            },
-            {
-                "category": "HARM_CATEGORY_HATE_SPEECH",
-                "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-            },
-            {
-                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-                "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-            },
-            {
-                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-                "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-            }
-        ]
-
-        # プロンプト構築
-        prompt = f"{system_instruction}\n\n{text}"
-
-        # コンテンツ生成
-        response = model.generate_content(
-            prompt,
-            generation_config=generation_config,
-            safety_settings=safety_settings
-        )
-
-        # レスポンスからHTMLを抽出
-        raw_response = response.text
-
-        # HTMLタグ部分だけを抽出（```html と ``` の間）
-        html_start = raw_response.find("```html")
-        html_end = raw_response.rfind("```")
+        logger.info(f"使用する Gemini モデル: {model_name}")
+
+        if model_name == "gemini-2.5-flash-preview-04-17":
+            # 新ライブラリ(genai_new)を使用し thinkingBudget=0 を設定
+            client = genai_new.Client(api_key=api_key)
+            logger.info("新ライブラリ genai_new を使用 (thinkingBudget=0)")
+            cfg = types.GenerateContentConfig(
+                thinking_config=types.ThinkingConfig(thinking_budget=0)
+            )
+            response = client.models.generate_content(
+                model=model_name,
+                contents=text,
+                config=cfg
+            )
+            raw = response.text
+        else:
+            # 既存ライブラリ(genai_old)のまま
+            genai_old.configure(api_key=api_key)
+            system_instruction = load_system_instruction(style)
+            prompt = f"{system_instruction}\n\n{text}"
+            response = genai_old.GenerativeModel(model_name).generate_content(
+                prompt,
+                generation_config={
+                    "temperature": temperature,
+                    "top_p": 0.7,
+                    "top_k": 20,
+                    "max_output_tokens": 8192,
+                    "candidate_count": 1
+                },
+                safety_settings=[
+                    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+                    {"category": "HARM_CATEGORY_HATE_SPEECH",      "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+                    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+                    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT","threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+                ]
+            )
+            raw = response.text
 
+        # Markdown ```html``` 部分を取り出す
+        html_start = raw.find("```html")
+        html_end   = raw.rfind("```")
         if html_start != -1 and html_end != -1 and html_start < html_end:
-            html_start += 7  # "```html" の長さ分進める
-            html_code = raw_response[html_start:html_end].strip()
-            logger.info(f"HTMLの生成に成功: 長さ = {len(html_code)}")
-
-            # Font Awesomeのレイアウト改善
-            html_code = enhance_font_awesome_layout(html_code)
-            logger.info("Font Awesomeレイアウトの最適化を適用しました")
-
-            return html_code
+            html_code = raw[html_start + 7:html_end].strip()
         else:
-            # HTMLタグが見つからない場合、レスポンス全体を返す
-            logger.warning("レスポンスから ```html ``` タグが見つかりませんでした。全テキストを返します。")
-            return raw_response
+            html_code = raw
+
+        # Font Awesome レイアウト最適化
+        html_code = enhance_font_awesome_layout(html_code)
+        logger.info("Font Awesome レイアウトの最適化を適用しました")
+        return html_code
 
     except Exception as e:
-        logger.error(f"HTML生成中にエラーが発生: {e}", exc_info=True)
-        raise Exception(f"Gemini APIでのHTML生成に失敗しました: {e}")
+        logger.error(f"HTML生成中にエラー: {e}", exc_info=True)
+        raise Exception(f"Gemini API での HTML 生成に失敗しました: {e}")
 
-# 画像から余分な空白領域をトリミングする関数 - NumPyを使って最適化
+# --- 画像トリミング ---
 def trim_image_whitespace(image, threshold=250, padding=10):
-    """
-    NumPyを使用して最適化された画像トリミング関数
-    
-    Args:
-        image: PIL.Image - 入力画像
-        threshold: int - どの明るさ以上を空白と判断するか (0-255)
-        padding: int - トリミング後に残す余白のピクセル数
-
-    Returns:
-        トリミングされたPIL.Image
-    """
     try:
-        # グレースケールに変換
         gray = image.convert('L')
-        
-        # NumPy配列として取得（高速処理のため）
         np_image = np.array(gray)
-        
-        # マスク作成（非白ピクセル）
         mask = np_image < threshold
-        
-        # マスクから行と列のインデックスを取得
         rows = np.any(mask, axis=1)
         cols = np.any(mask, axis=0)
-        
-        # 非空のインデックス範囲を取得
         if np.any(rows) and np.any(cols):
             row_indices = np.where(rows)[0]
             col_indices = np.where(cols)[0]
-            
-            # 範囲取得
             min_y, max_y = row_indices[0], row_indices[-1]
             min_x, max_x = col_indices[0], col_indices[-1]
-            
-            # パディング追加
             min_x = max(0, min_x - padding)
             min_y = max(0, min_y - padding)
             max_x = min(image.width - 1, max_x + padding)
             max_y = min(image.height - 1, max_y + padding)
-            
-            # 画像をトリミング
             trimmed = image.crop((min_x, min_y, max_x + 1, max_y + 1))
-            
-            logger.info(f"画像をトリミングしました: 元サイズ {image.width}x{image.height} → トリミング後 {trimmed.width}x{trimmed.height}")
+            logger.info(f"画像をトリミングしました: {image.width}x{image.height} → {trimmed.width}x{trimmed.height}")
             return trimmed
-        
         logger.warning("トリミング領域が見つかりません。元の画像を返します。")
         return image
-        
     except Exception as e:
         logger.error(f"画像トリミング中にエラー: {e}", exc_info=True)
-        return image  # エラー時は元の画像を返す
+        return image
 
-# 最適化されたスクリーンショット生成関数 - 外部から初期化済みドライバーを受け取れるように
+# --- スクリーンショット生成 ---
 def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0,
                               trim_whitespace: bool = True, driver=None) -> Image.Image:
-    """
-    Renders HTML code to a full-page screenshot using Selenium.
-    Optimized to accept an external driver or get one from the pool.
-
-    Args:
-        html_code: The HTML source code string.
-        extension_percentage: Percentage of extra space to add vertically.
-        trim_whitespace: Whether to trim excess whitespace from the image.
-        driver: An optional pre-initialized WebDriver instance.
-
-    Returns:
-        A PIL Image object of the screenshot.
-    """
     tmp_path = None
     driver_from_pool = False
-    
-    # ドライバーがない場合はプールから取得
     if driver is None:
         driver = driver_pool.get_driver()
         driver_from_pool = True
         logger.info("WebDriverプールからドライバーを取得しました")
 
-    # 1) Save HTML code to a temporary file
     try:
+        # HTML を一時ファイルに保存
         with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode='w', encoding='utf-8') as tmp_file:
             tmp_path = tmp_file.name
             tmp_file.write(html_code)
         logger.info(f"HTML saved to temporary file: {tmp_path}")
-    except Exception as e:
-        logger.error(f"Error writing temporary HTML file: {e}")
-        if driver_from_pool:
-            driver_pool.release_driver(driver)
-        return Image.new('RGB', (1, 1), color=(0, 0, 0))
 
-    try:
-        # ウィンドウサイズ初期設定
+        # ウィンドウ初期サイズ設定
         initial_width = 1200
         initial_height = 1000
         driver.set_window_size(initial_width, initial_height)
-        file_url = "file://" + tmp_path
-        logger.info(f"Navigating to {file_url}")
-        driver.get(file_url)
+        driver.get("file://" + tmp_path)
 
-        # ページ読み込み待機 - 動的な待機時間を実装
-        logger.info("Waiting for body element...")
+        # body 要素の読み込み待機
         WebDriverWait(driver, 10).until(
             EC.presence_of_element_located((By.TAG_NAME, "body"))
         )
-        logger.info("Body element found. Waiting for resource loading...")
 
-        # リソース読み込みの動的待機 - 最適化
-        max_wait = 5  # 最大待機時間（秒）
-        wait_increment = 0.2  # 確認間隔
+        # リソース読み込み待機
+        max_wait = 5
+        wait_increment = 0.2
         wait_time = 0
-        
         while wait_time < max_wait:
-            resource_state = driver.execute_script("""
+            state = driver.execute_script("""
                 return {
                     complete: document.readyState === 'complete',
                     imgCount: document.images.length,
@@ -482,48 +357,34 @@ def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0
                     faElements: document.querySelectorAll('.fa, .fas, .far, .fab, [class*="fa-"]').length
                 };
             """)
-            
-            # ドキュメント完了かつ画像が読み込まれている場合、待機終了
-            if resource_state['complete'] and (resource_state['imgCount'] == 0 or 
-                                             resource_state['imgLoaded'] == resource_state['imgCount']):
-                logger.info(f"リソース読み込み完了: {resource_state}")
+            if state['complete'] and (state['imgCount'] == 0 or state['imgLoaded'] == state['imgCount']):
                 break
-                
             time.sleep(wait_increment)
             wait_time += wait_increment
-            logger.info(f"リソース待機中... {wait_time:.1f}秒経過, 状態: {resource_state}")
-        
-        # Font Awesome要素が多い場合は追加待機
-        fa_count = resource_state.get('faElements', 0)
-        if fa_count > 30:
-            logger.info(f"{fa_count}個のFont Awesome要素があるため、追加待機...")
-            time.sleep(min(1.0, fa_count / 100))  # 要素数に応じて待機（最大1秒）
-
-        # コンテンツレンダリングのためのスクロール処理 - パフォーマンス改善
-        logger.info("Performing content rendering scroll...")
+
+        # Font Awesome 要素��多い場合は少し待機
+        if state.get('faElements', 0) > 30:
+            time.sleep(min(1.0, state['faElements'] / 100))
+
+        # スクロール処理
         total_height = driver.execute_script("return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight);")
         viewport_height = driver.execute_script("return window.innerHeight;")
-        scrolls_needed = max(1, min(5, total_height // viewport_height))  # 最大5回までに制限
-        
-        # スクロール処理の高速化
+        scrolls_needed = max(1, min(5, total_height // viewport_height))
         for i in range(scrolls_needed):
-            scroll_pos = i * (viewport_height - 100)  # 少しだけオーバーラップ
+            scroll_pos = i * (viewport_height - 100)
             driver.execute_script(f"window.scrollTo(0, {scroll_pos});")
-            time.sleep(0.1)  # 高速化のため待機時間短縮
-        
-        # トップに戻る
+            time.sleep(0.1)
         driver.execute_script("window.scrollTo(0, 0);")
-        time.sleep(0.2)  # 短い待機に変更
-        logger.info("Scroll rendering completed")
+        time.sleep(0.2)
 
-        # スクロールバーを非表示に
+        # スクロールバー非表示
         driver.execute_script("""
             document.documentElement.style.overflow = 'hidden';
             document.body.style.overflow = 'hidden';
         """)
-        
-        # ページの寸法を取得
-        dimensions = driver.execute_script("""
+
+        # ページ寸法取得
+        dims = driver.execute_script("""
             return {
                 width: Math.max(
                     document.documentElement.scrollWidth,
@@ -543,39 +404,23 @@ def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0
                 )
             };
         """)
-        scroll_width = dimensions['width']
-        scroll_height = dimensions['height']
-        logger.info(f"Detected dimensions: width={scroll_width}, height={scroll_height}")
-
-        # 最小/最大値の設定
-        scroll_width = max(scroll_width, 100)
-        scroll_height = max(scroll_height, 100)
+        scroll_width = max(dims['width'], 100)
+        scroll_height = max(dims['height'], 100)
         scroll_width = min(scroll_width, 2000)
         scroll_height = min(scroll_height, 4000)
-        
-        # レイアウト安定化のための待機 
-        time.sleep(2.0)  
 
-        # 高さに余白を追加
+        # 縦余白追加
         adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0))
         adjusted_height = max(adjusted_height, scroll_height, 100)
-        
-        # ウィンドウサイズを調整
-        adjusted_width = scroll_width
-        logger.info(f"Resizing window to: width={adjusted_width}, height={adjusted_height}")
-        driver.set_window_size(adjusted_width, adjusted_height)
-        time.sleep(0.5)  # 短縮した待機時間
+        driver.set_window_size(scroll_width, adjusted_height)
+        time.sleep(0.5)
 
         # スクリーンショット取得
-        logger.info("Taking screenshot...")
         png = driver.get_screenshot_as_png()
-        logger.info("Screenshot taken successfully.")
-
-        # PIL画像に変換
         img = Image.open(BytesIO(png))
         logger.info(f"Screenshot dimensions: {img.width}x{img.height}")
 
-        # 余白トリミング - 最適化版を使用
+        # 余白トリミング
         if trim_whitespace:
             img = trim_image_whitespace(img, threshold=248, padding=20)
             logger.info(f"Trimmed dimensions: {img.width}x{img.height}")
@@ -584,15 +429,11 @@ def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0
 
     except Exception as e:
         logger.error(f"Error during screenshot generation: {e}", exc_info=True)
-        # エラー時は小さい黒画像を返す
         return Image.new('RGB', (1, 1), color=(0, 0, 0))
+
     finally:
-        logger.info("Cleaning up...")
-        # WebDriverプールに戻す
         if driver_from_pool:
             driver_pool.release_driver(driver)
-            logger.info("Returned driver to pool")
-        # 一時ファイル削除
         if tmp_path and os.path.exists(tmp_path):
             try:
                 os.remove(tmp_path)
@@ -600,188 +441,128 @@ def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0
             except Exception as e:
                 logger.error(f"Error removing temporary file {tmp_path}: {e}")
 
-# --- 並列処理を活用した新しい関数 ---
+# --- 並列処理版スクリーンショット生成 ---
 def text_to_screenshot_parallel(text: str, extension_percentage: float, temperature: float = 0.5,
-                             trim_whitespace: bool = True, style: str = "standard") -> Image.Image:
-    """テキストをGemini APIでHTMLに変換し、並列処理でスクリーンショットを生成する関数"""
+                                trim_whitespace: bool = True, style: str = "standard") -> Image.Image:
     start_time = time.time()
-    logger.info("並列処理によるテキスト→スクリーンショット生成を開始")
-    
+    drv = None
+    tmp_path = None
+    driver_from_pool = False
     try:
-        # WebDriverと HTML生成を並列で実行
         with ThreadPoolExecutor(max_workers=2) as executor:
-            # Gemini APIリクエストタスク
-            html_future = executor.submit(
-                generate_html_from_text, 
-                text=text, 
-                temperature=temperature, 
-                style=style
-            )
-            
-            # WebDriver初期化タスク - プール使用
+            html_future = executor.submit(generate_html_from_text, text, temperature, style)
             driver_future = executor.submit(driver_pool.get_driver)
-            
-            # 結果を取得
             html_code = html_future.result()
-            driver = driver_future.result()
-            
-        # ドライバーはプールから取得しているためフラグ設定
-        driver_from_pool = True
-        
-        # HTMLコードとドライバーが準備できたらスクリーンショット生成
-        logger.info(f"HTML生成完了：{len(html_code)}文字。スクリーンショット生成開始。")
-        
-        # レンダリング前にドライバーの初期設定
-        tmp_path = None
-        try:
-            # 一時ファイルにHTMLを保存
-            with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode='w', encoding='utf-8') as tmp_file:
-                tmp_path = tmp_file.name
-                tmp_file.write(html_code)
-            logger.info(f"HTMLを一時ファイルに保存: {tmp_path}")
-            
-            # ウィンドウサイズ初期設定
-            initial_width = 1200
-            initial_height = 1000
-            driver.set_window_size(initial_width, initial_height)
-            file_url = "file://" + tmp_path
-            logger.info(f"ページに移動: {file_url}")
-            driver.get(file_url)
-            
-            # ここからスクリーンショット生成ロジック（前の実装と同様）
-            # ページ読み込み待機 - 動的な待機時間を実装
-            logger.info("body要素を待機...")
-            WebDriverWait(driver, 10).until(
-                EC.presence_of_element_located((By.TAG_NAME, "body"))
-            )
-            logger.info("body要素を検出。リソース読み込みを待機...")
-            
-            # リソース読み込みの動的待機 - 最適化
-            max_wait = 3  # 最大待機時間（秒）
-            wait_increment = 0.2  # 確認間隔
-            wait_time = 0
-            
-            while wait_time < max_wait:
-                resource_state = driver.execute_script("""
-                    return {
-                        complete: document.readyState === 'complete',
-                        imgCount: document.images.length,
-                        imgLoaded: Array.from(document.images).filter(img => img.complete).length,
-                        faElements: document.querySelectorAll('.fa, .fas, .far, .fab, [class*="fa-"]').length
-                    };
-                """)
-                
-                # ドキュメント完了かつ画像が読み込まれている場合、待機終了
-                if resource_state['complete'] and (resource_state['imgCount'] == 0 or 
-                                                 resource_state['imgLoaded'] == resource_state['imgCount']):
-                    logger.info(f"リソース読み込み完了: {resource_state}")
-                    break
-                    
-                time.sleep(wait_increment)
-                wait_time += wait_increment
-            
-            # Font Awesome要素が多い場合は追加待機
-            fa_count = resource_state.get('faElements', 0)
-            if fa_count > 30:
-                logger.info(f"{fa_count}個のFont Awesome要素があるため、追加待機...")
-                time.sleep(min(1.0, fa_count / 100))  # 要素数に応じて待機（最大0.5秒）
-            
-            # コンテンツレンダリングのための簡易スクロール
-            driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
-            time.sleep(0.2)
-            driver.execute_script("window.scrollTo(0, 0);")
-            time.sleep(0.2)
-            
-            # スクロールバーを非表示に
-            driver.execute_script("""
-                document.documentElement.style.overflow = 'hidden';
-                document.body.style.overflow = 'hidden';
-            """)
-            
-            # ページの寸法を取得
-            dimensions = driver.execute_script("""
+            drv = driver_future.result()
+            driver_from_pool = True
+
+        # HTML→一時ファイル
+        with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode='w', encoding='utf-8') as tmp_file:
+            tmp_path = tmp_file.name
+            tmp_file.write(html_code)
+        logger.info(f"HTMLを一時ファイルに保存: {tmp_path}")
+
+        # ドライバ初期化
+        drv.set_window_size(1200, 1000)
+        drv.get("file://" + tmp_path)
+        WebDriverWait(drv, 10).until(EC.presence_of_element_located((By.TAG_NAME, "body")))
+
+        # リソース待機
+        max_wait = 3
+        wait_increment = 0.2
+        wait_time = 0
+        while wait_time < max_wait:
+            state = drv.execute_script("""
                 return {
-                    width: Math.max(
-                        document.documentElement.scrollWidth,
-                        document.documentElement.offsetWidth,
-                        document.documentElement.clientWidth,
-                        document.body ? document.body.scrollWidth : 0,
-                        document.body ? document.body.offsetWidth : 0,
-                        document.body ? document.body.clientWidth : 0
-                    ),
-                    height: Math.max(
-                        document.documentElement.scrollHeight,
-                        document.documentElement.offsetHeight,
-                        document.documentElement.clientHeight,
-                        document.body ? document.body.scrollHeight : 0,
-                        document.body ? document.body.offsetHeight : 0,
-                        document.body ? document.body.clientHeight : 0
-                    )
+                    complete: document.readyState==='complete',
+                    imgCount: document.images.length,
+                    imgLoaded: Array.from(document.images).filter(img=>img.complete).length,
+                    faElements: document.querySelectorAll('.fa, .fas, .far, .fab, [class*="fa-"]').length
                 };
             """)
-            scroll_width = dimensions['width']
-            scroll_height = dimensions['height']
-            
-            # 最小/最大値の設定
-            scroll_width = max(scroll_width, 100)
-            scroll_height = max(scroll_height, 100)
-            scroll_width = min(scroll_width, 2000)
-            scroll_height = min(scroll_height, 4000)
-            
-            # 高さに余白を追加
-            adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0))
-            adjusted_height = max(adjusted_height, scroll_height, 100)
-            
-            # ウィンドウサイズを調整
-            driver.set_window_size(scroll_width, adjusted_height)
-            time.sleep(0.2)
-            
-            # スクリーンショット取得
-            logger.info("スクリーンショットを撮影...")
-            png = driver.get_screenshot_as_png()
-            
-            # PIL画像に変換
-            img = Image.open(BytesIO(png))
-            logger.info(f"スクリーンショットサイズ: {img.width}x{img.height}")
-            
-            # 余白トリミング
-            if trim_whitespace:
-                img = trim_image_whitespace(img, threshold=248, padding=20)
-                logger.info(f"トリミング後のサイズ: {img.width}x{img.height}")
-            
-            elapsed = time.time() - start_time
-            logger.info(f"並列処理による生成完了。所要時間: {elapsed:.2f}秒")
-            return img
-            
-        except Exception as e:
-            logger.error(f"スクリーンショット生成中にエラー: {e}", exc_info=True)
-            return Image.new('RGB', (1, 1), color=(0, 0, 0))
-        finally:
-            # WebDriverプールに戻す
-            if driver_from_pool:
-                driver_pool.release_driver(driver)
-            # 一時ファイル削除
-            if tmp_path and os.path.exists(tmp_path):
-                try:
-                    os.remove(tmp_path)
-                except Exception as e:
-                    logger.error(f"一時ファイル削除エラー: {e}")
-    
+            if state['complete'] and (state['imgCount']==0 or state['imgLoaded']==state['imgCount']):
+                break
+            time.sleep(wait_increment)
+            wait_time += wait_increment
+
+        if state.get('faElements', 0) > 30:
+            time.sleep(min(1.0, state['faElements'] / 100))
+
+        # 簡易スクロール
+        drv.execute_script("window.scrollTo(0, document.body.scrollHeight);")
+        time.sleep(0.2)
+        drv.execute_script("window.scrollTo(0, 0);")
+        time.sleep(0.2)
+        drv.execute_script("document.documentElement.style.overflow='hidden';document.body.style.overflow='hidden';")
+
+        # 寸法取得
+        dims = drv.execute_script("""
+            return {
+                width: Math.max(
+                    document.documentElement.scrollWidth,
+                    document.documentElement.offsetWidth,
+                    document.documentElement.clientWidth,
+                    document.body ? document.body.scrollWidth : 0,
+                    document.body ? document.body.offsetWidth : 0,
+                    document.body ? document.body.clientWidth : 0
+                ),
+                height: Math.max(
+                    document.documentElement.scrollHeight,
+                    document.documentElement.offsetHeight,
+                    document.documentElement.clientHeight,
+                    document.body ? document.body.scrollHeight : 0,
+                    document.body ? document.body.offsetHeight : 0,
+                    document.body ? document.body.clientHeight : 0
+                )
+            };
+        """)
+        w = max(dims['width'], 100)
+        h = max(dims['height'], 100)
+        w = min(w, 2000)
+        h = min(h, 4000)
+        adjusted_h = int(h * (1 + extension_percentage / 100.0))
+        adjusted_h = max(adjusted_h, h, 100)
+
+        drv.set_window_size(w, adjusted_h)
+        time.sleep(0.2)
+
+        # スクリーンショット取得
+        png = drv.get_screenshot_as_png()
+        img = Image.open(BytesIO(png))
+        if trim_whitespace:
+            img = trim_image_whitespace(img, threshold=248, padding=20)
+
+        elapsed = time.time() - start_time
+        logger.info(f"Parallel generation 完了 (所要時間: {elapsed:.2f}秒)")
+        return img
+
     except Exception as e:
-        logger.error(f"並列処理中のエラー: {e}", exc_info=True)
-        return Image.new('RGB', (1, 1), color=(0, 0, 0))  # エラー時は黒画像
+        logger.error(f"Parallel generation error: {e}", exc_info=True)
+        return Image.new('RGB', (1, 1), color=(0, 0, 0))
+
+    finally:
+        if driver_from_pool and drv:
+            driver_pool.release_driver(drv)
+        if tmp_path and os.path.exists(tmp_path):
+            try:
+                os.remove(tmp_path)
+            except Exception:
+                pass
 
-# 従来の非並列版も残す（互換性のため）
+# --- レガシー版スクリーンショット生成 ---
 def text_to_screenshot(text: str, extension_percentage: float, temperature: float = 0.3,
-                    trim_whitespace: bool = True, style: str = "standard") -> Image.Image:
-    """テキストをGemini APIでHTMLに変換し、スクリーンショットを生成する統合関数（レガシー版）"""
-    # 並列処理版を呼び出す
+                       trim_whitespace: bool = True, style: str = "standard") -> Image.Image:
     return text_to_screenshot_parallel(text, extension_percentage, temperature, trim_whitespace, style)
 
+# --- 入力モード切り替え用関数 ---
+def process_input(input_mode, input_text, extension_percentage, temperature, trim_whitespace, style):
+    if input_mode == "HTML入力":
+        return render_fullpage_screenshot(input_text, extension_percentage, trim_whitespace)
+    else:
+        return text_to_screenshot_parallel(input_text, extension_percentage, temperature, trim_whitespace, style)
+
 # --- FastAPI Setup ---
 app = FastAPI()
-
-# CORS設定を追加
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -791,88 +572,58 @@ app.add_middleware(
 )
 
 # 静的ファイルのサービング設定
-# Gradioのディレクトリを探索してアセットを見つける
 gradio_dir = os.path.dirname(gr.__file__)
-logger.info(f"Gradio version: {gr.__version__}")
-logger.info(f"Gradio directory: {gradio_dir}")
-
-# 基本的な静的ファイルディレクトリをマウント
 static_dir = os.path.join(gradio_dir, "templates", "frontend", "static")
 if os.path.exists(static_dir):
-    logger.info(f"Mounting static directory: {static_dir}")
     app.mount("/static", StaticFiles(directory=static_dir), name="static")
-
-# _appディレクトリを探す（新しいSvelteKitベースのフロントエンド用）
 app_dir = os.path.join(gradio_dir, "templates", "frontend", "_app")
 if os.path.exists(app_dir):
-    logger.info(f"Mounting _app directory: {app_dir}")
     app.mount("/_app", StaticFiles(directory=app_dir), name="_app")
-
-# assetsディレクトリを探す
 assets_dir = os.path.join(gradio_dir, "templates", "frontend", "assets")
 if os.path.exists(assets_dir):
-    logger.info(f"Mounting assets directory: {assets_dir}")
     app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")
-
-# cdnディレクトリがあれば追加
 cdn_dir = os.path.join(gradio_dir, "templates", "cdn")
 if os.path.exists(cdn_dir):
-    logger.info(f"Mounting cdn directory: {cdn_dir}")
     app.mount("/cdn", StaticFiles(directory=cdn_dir), name="cdn")
 
-
-# API Endpoint for screenshot generation
+# --- API Endpoint for HTML→Screenshot ---
 @app.post("/api/screenshot",
           response_class=StreamingResponse,
           tags=["Screenshot"],
           summary="Render HTML to Full Page Screenshot",
           description="Takes HTML code and an optional vertical extension percentage, renders it using a headless browser, and returns the full-page screenshot as a PNG image.")
 async def api_render_screenshot(request: ScreenshotRequest):
-    """
-    API endpoint to render HTML and return a screenshot.
-    """
     try:
         logger.info(f"API request received. Extension: {request.extension_percentage}%")
-        # Run the blocking Selenium code (now using the pooled version)
         pil_image = render_fullpage_screenshot(
             request.html_code,
             request.extension_percentage,
             request.trim_whitespace
         )
-
         if pil_image.size == (1, 1):
-             logger.error("Screenshot generation failed, returning 1x1 error image.")
-             # Optionally return a proper error response instead of 1x1 image
-             # raise HTTPException(status_code=500, detail="Failed to generate screenshot")
-
-        # Convert PIL Image to PNG bytes
+            logger.error("Screenshot generation failed, returning 1x1 error image.")
         img_byte_arr = BytesIO()
         pil_image.save(img_byte_arr, format='PNG')
-        img_byte_arr.seek(0)  # Go to the start of the BytesIO buffer
-
+        img_byte_arr.seek(0)
         logger.info("Returning screenshot as PNG stream.")
         return StreamingResponse(img_byte_arr, media_type="image/png")
-
     except Exception as e:
         logger.error(f"API Error: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")
 
-# --- 新しいGemini API連携エンドポイント（並列処理版） ---
+# --- API Endpoint for Text→Infographic Screenshot ---
 @app.post("/api/text-to-screenshot",
           response_class=StreamingResponse,
           tags=["Screenshot", "Gemini"],
           summary="テキストからインフォグラフィックを生成",
           description="テキストをGemini APIを使ってHTMLインフォグラフィックに変換し、スクリーンショットとして返します。")
 async def api_text_to_screenshot(request: GeminiRequest):
-    """
-    テキストからHTMLインフォグラフィックを生成してスクリーンショットを返すAPIエンドポイント
-    """
     try:
-        logger.info(f"テキスト→スクリーンショットAPIリクエスト受信。テキスト長さ: {len(request.text)}, "
-                   f"拡張率: {request.extension_percentage}%, 温度: {request.temperature}, "
-                   f"スタイル: {request.style}")
-
-        # 並列処理版を使用
+        logger.info(
+            f"テキスト→スクリーンショットAPIリクエスト受信。"
+            f"テキスト長さ: {len(request.text)}, 拡張率: {request.extension_percentage}%, "
+            f"温度: {request.temperature}, スタイル: {request.style}"
+        )
         pil_image = text_to_screenshot_parallel(
             request.text,
             request.extension_percentage,
@@ -880,34 +631,18 @@ async def api_text_to_screenshot(request: GeminiRequest):
             request.trim_whitespace,
             request.style
         )
-
         if pil_image.size == (1, 1):
             logger.error("スクリーンショット生成に失敗しました。1x1エラー画像を返します。")
-
-        # PIL画像をPNGバイトに変換
         img_byte_arr = BytesIO()
         pil_image.save(img_byte_arr, format='PNG')
-        img_byte_arr.seek(0)  # BytesIOバッファの先頭に戻る
-
+        img_byte_arr.seek(0)
         logger.info("スクリーンショットをPNGストリームとして返します。")
         return StreamingResponse(img_byte_arr, media_type="image/png")
-
     except Exception as e:
         logger.error(f"API Error: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")
 
 # --- Gradio Interface Definition ---
-# 入力モードの選択用Radioコンポーネント
-def process_input(input_mode, input_text, extension_percentage, temperature, trim_whitespace, style):
-    """入力モードに応じて適切な処理を行う"""
-    if input_mode == "HTML入力":
-        # HTMLモードの場合は既存の処理（スタイルは使わない）
-        return render_fullpage_screenshot(input_text, extension_percentage, trim_whitespace)
-    else:
-        # テキスト入力モードの場合はGemini APIを使用（並列処理版）
-        return text_to_screenshot_parallel(input_text, extension_percentage, temperature, trim_whitespace, style)
-
-# Gradio UIの定義
 with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr.themes.Base()) as iface:
     gr.Markdown("# HTMLビューア & テキスト→インフォグラフィック変換")
     gr.Markdown("HTMLコードをレンダリングするか、テキストをGemini APIでインフォグラフィックに変換して画像として取得します。")
@@ -920,7 +655,6 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
             value="HTML入力"
         )
 
-    # 共用のテキストボックス
     input_text = gr.Textbox(
         lines=15,
         label="入力",
@@ -929,35 +663,30 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
 
     with gr.Row():
         with gr.Column(scale=1):
-            # スタイル選択ドロップダウン
             style_dropdown = gr.Dropdown(
                 choices=["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"],
                 value="standard",
                 label="デザインスタイル",
                 info="テキスト→HTML変換時のデザインテーマを選択します",
-                visible=False  # テキスト入力モードの時だけ表示
+                visible=False
             )
-
         with gr.Column(scale=2):
             extension_percentage = gr.Slider(
                 minimum=0,
                 maximum=30,
                 step=1.0,
-                value=10,    # デフォルト値10%
+                value=10,
                 label="上下高さ拡張率（%）"
             )
-
-            # 温度調整スライダー（テキストモード時のみ表示）
             temperature = gr.Slider(
                 minimum=0.0,
                 maximum=1.0,
                 step=0.1,
-                value=0.5,    # デフォルト値を0.5に設定
+                value=0.5,
                 label="生成時の温度（低い=一貫性高、高い=創造性高）",
-                visible=False  # 最初は非表示
+                visible=False
             )
 
-    # 余白トリミングオプション
     trim_whitespace = gr.Checkbox(
         label="余白を自動トリミング",
         value=True,
@@ -967,10 +696,8 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
     submit_btn = gr.Button("生成")
     output_image = gr.Image(type="pil", label="ページ全体のスクリーンショット")
 
-    # 入力モード変更時のイベント処理（テキストモード時のみ温度スライダーとスタイルドロップダウンを表示）
     def update_controls_visibility(mode):
-        # Gradio 4.x用のアップデート方法
-        is_text_mode = mode == "テキスト入力"
+        is_text_mode = (mode == "テキスト入力")
         return [
             gr.update(visible=is_text_mode),  # temperature
             gr.update(visible=is_text_mode),  # style_dropdown
@@ -982,14 +709,12 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
         outputs=[temperature, style_dropdown]
     )
 
-    # 生成ボタンクリック時のイベント処理
     submit_btn.click(
         fn=process_input,
         inputs=[input_mode, input_text, extension_percentage, temperature, trim_whitespace, style_dropdown],
         outputs=output_image
     )
 
-    # 環境変数情報を表示
     gemini_model = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro")
     gr.Markdown(f"""
     ## APIエンドポイント
@@ -1005,7 +730,7 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
 # --- Mount Gradio App onto FastAPI ---
 app = gr.mount_gradio_app(app, iface, path="/")
 
-# --- Run with Uvicorn (for local testing) ---
+# --- ローカル開発用 Uvicorn 起動 ---
 if __name__ == "__main__":
     import uvicorn
     logger.info("Starting Uvicorn server for local development...")
@@ -1013,4 +738,4 @@ if __name__ == "__main__":
 
 # アプリケーション終了時にWebDriverプールをクリーンアップ
 import atexit
-atexit.register(driver_pool.close_all)
+atexit.register(driver_pool.close_all)