Update app.py

app.py

@@ -1,11 +1,9 @@
-# app.py  (Gradio 5.x 対応フルバージョン)
-# =========================================
-# 主要変更点
-#   • gradio==4.19.2  → gradio>=5.29.0   （requirements.txt で指定）
-#   • Blocks の theme を Origin() に変更（4.x の外観を継承）
-#   • mount_gradio_app に ssr_mode=False を追加
-#   • Gradio 内部アセットの手動マウントを撤去（5.x では不要）
-# =========================================
+# ===============================================================
+# app.py  ―  Gradio 5.x + FastAPI + Gemini + Selenium
+# ・Gradio 5.29.0 以上を前提に最適化
+# ・UI/API を維持したまま 307 ループを解消
+# ・/gradio にサブマウントし / はリダイレクトのみ
+# ===============================================================
 
 import os
 import time
@@ -16,12 +14,12 @@ import queue
 from io import BytesIO
 from concurrent.futures import ThreadPoolExecutor
 
-import numpy as np                   # 画像トリミング高速化
+import numpy as np
 from PIL import Image
 
-import gradio as gr                  # ★ 5.x
+import gradio as gr
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, RedirectResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
@@ -31,55 +29,61 @@ from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 
-import google.generativeai as genai  # Gemini
+import google.generativeai as genai
 from huggingface_hub import hf_hub_download
 
-# -------------------- ロガー設定 --------------------
+
+# ===============================================================
+# ロガー
+# ===============================================================
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ====================================================
-# WebDriver プール実装
-# ====================================================
+
+# ===============================================================
+# WebDriver プール
+# ===============================================================
 class WebDriverPool:
+    """複数の WebDriver を使い回すシンプルなプール"""
     def __init__(self, max_drivers: int = 3):
-        self.driver_queue = queue.Queue()
+        self.driver_queue: "queue.Queue[webdriver.Chrome]" = queue.Queue()
         self.max_drivers = max_drivers
         self.lock = threading.Lock()
         self.count = 0
         logger.info(f"WebDriver プール初期化: 最大 {max_drivers}")
 
-    def get_driver(self):
-        # 既存
+    def _create_driver(self) -> webdriver.Chrome:
+        options = Options()
+        options.add_argument("--headless")
+        options.add_argument("--no-sandbox")
+        options.add_argument("--disable-dev-shm-usage")
+        options.add_argument("--force-device-scale-factor=1")
+        options.add_argument("--disable-features=NetworkService")
+        options.add_argument("--dns-prefetch-disable")
+
+        chromedriver_path = os.environ.get("CHROMEDRIVER_PATH")
+        if chromedriver_path and os.path.exists(chromedriver_path):
+            logger.info(f"環境変数 CHROMEDRIVER_PATH を使用: {chromedriver_path}")
+            service = webdriver.ChromeService(executable_path=chromedriver_path)
+            return webdriver.Chrome(service=service, options=options)
+        return webdriver.Chrome(options=options)
+
+    def get_driver(self) -> webdriver.Chrome:
         if not self.driver_queue.empty():
-            logger.info("既存 WebDriver を取得")
+            logger.info("プールから既存 WebDriver を取得")
             return self.driver_queue.get()
 
-        # 新規
         with self.lock:
             if self.count < self.max_drivers:
                 self.count += 1
                 logger.info(f"新規 WebDriver 作成 ({self.count}/{self.max_drivers})")
-                options = Options()
-                options.add_argument("--headless")
-                options.add_argument("--no-sandbox")
-                options.add_argument("--disable-dev-shm-usage")
-                options.add_argument("--force-device-scale-factor=1")
-                options.add_argument("--disable-features=NetworkService")
-                options.add_argument("--dns-prefetch-disable")
-
-                driver_path = os.environ.get("CHROMEDRIVER_PATH")
-                if driver_path and os.path.exists(driver_path):
-                    logger.info(f"環境変数 CHROMEDRIVER_PATH 使用: {driver_path}")
-                    service = webdriver.ChromeService(executable_path=driver_path)
-                    return webdriver.Chrome(service=service, options=options)
-                return webdriver.Chrome(options=options)
-
-        # プール満杯
-        logger.info("プール満杯。返却待ち…")
+                return self._create_driver()
+
+        # ここに来るのはプール満杯時
+        logger.info("WebDriver プール満杯。空きを待機")
         return self.driver_queue.get()
 
-    def release_driver(self, driver):
+    def release_driver(self, driver: webdriver.Chrome):
         if driver:
             try:
                 driver.get("about:blank")
@@ -90,18 +94,18 @@ class WebDriverPool:
                 self.driver_queue.put(driver)
                 logger.info("WebDriver をプールに返却")
             except Exception as e:
-                logger.error(f"返却時エラー: {e}")
+                logger.error(f"返却エラー: {e}")
                 driver.quit()
                 with self.lock:
                     self.count -= 1
 
     def close_all(self):
-        logger.info("全 WebDriver 終了処理")
+        logger.info("WebDriver 全終了")
         closed = 0
         while not self.driver_queue.empty():
             try:
-                driver = self.driver_queue.get(block=False)
-                driver.quit()
+                drv = self.driver_queue.get(block=False)
+                drv.quit()
                 closed += 1
             except queue.Empty:
                 break
@@ -111,12 +115,14 @@ class WebDriverPool:
             self.count = 0
         logger.info(f"{closed} 個の WebDriver を終了")
 
+
 # グローバルプール
 driver_pool = WebDriverPool(max_drivers=int(os.environ.get("MAX_WEBDRIVERS", "3")))
 
-# ====================================================
+
+# ===============================================================
 # Pydantic モデル
-# ====================================================
+# ===============================================================
 class GeminiRequest(BaseModel):
     text: str
     extension_percentage: float = 10.0
@@ -124,17 +130,19 @@ class GeminiRequest(BaseModel):
     trim_whitespace: bool = True
     style: str = "standard"
 
+
 class ScreenshotRequest(BaseModel):
     html_code: str
     extension_percentage: float = 10.0
     trim_whitespace: bool = True
     style: str = "standard"
 
-# ====================================================
+
+# ===============================================================
 # 補助関数
-# ====================================================
+# ===============================================================
 def enhance_font_awesome_layout(html_code: str) -> str:
-    """Font Awesome レイアウトを調整し preload タグを付与"""
+    """Font Awesome をプリロード + レイアウト微調整 CSS を注入"""
     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>
@@ -162,8 +170,9 @@ def enhance_font_awesome_layout(html_code: str) -> str:
             return html_code[:body_start] + f'<head>{fa_preload}{fa_css}</head>' + html_code[body_start:]
     return f'<html><head>{fa_preload}{fa_css}</head>{html_code}</html>'
 
+
 def load_system_instruction(style: str = "standard") -> str:
-    """テーマ別 prompt.txt を読み込み"""
+    """スタイル別の prompt.txt をローカル or HF から読み込み"""
     valid = ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"]
     if style not in valid:
         logger.warning(f"無効 style '{style}' → 'standard'")
@@ -183,7 +192,7 @@ def load_system_instruction(style: str = "standard") -> str:
         with open(file_path, encoding="utf-8") as f:
             return f.read()
     except Exception as e:
-        logger.warning(f"HuggingFace 取得失敗 ({e}) → デフォルト prompt.txt")
+        logger.warning(f"HF 取得失敗 ({e}) → デフォルト prompt.txt")
         file_path = hf_hub_download(
             repo_id="tomo2chin2/GURAREKOstlyle",
             filename="prompt.txt",
@@ -192,21 +201,23 @@ def load_system_instruction(style: str = "standard") -> str:
         with open(file_path, encoding="utf-8") as f:
             return f.read()
 
+
 def generate_html_from_text(text: str, temperature: float = 0.5, style: str = "standard") -> str:
-    """Gemini で HTML 生成"""
+    """Gemini で与えられたテキストを HTML に整形して返す"""
     api_key = os.environ.get("GEMINI_API_KEY")
     if not api_key:
         raise ValueError("環境変数 GEMINI_API_KEY が未設定")
-
     model_name = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro")
+
     genai.configure(api_key=api_key)
+    model = genai.GenerativeModel(model_name)
 
     system_instruction = load_system_instruction(style)
-    model = genai.GenerativeModel(model_name)
+    prompt = f"{system_instruction}\n\n{text}"
 
     generation_config = dict(
-        temperature=temperature, top_p=0.7, top_k=20, max_output_tokens=8192,
-        candidate_count=1
+        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"},
@@ -214,67 +225,62 @@ def generate_html_from_text(text: str, temperature: float = 0.5, style: str = "s
         {"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)
-    raw = response.text
+    resp = model.generate_content(prompt, generation_config=generation_config, safety_settings=safety_settings)
+    raw = resp.text
 
-    html_start = raw.find("```html")
-    html_end = raw.rfind("```")
-    if html_start != -1 and html_end != -1 and html_start < html_end:
-        html_code = raw[html_start + 7:html_end].strip()
-    else:
-        html_code = raw
+    # ```html … ``` 抜き取り
+    start = raw.find("```html")
+    end = raw.rfind("```")
+    html_code = raw[start + 7:end].strip() if start != -1 and end != -1 else raw
+    return enhance_font_awesome_layout(html_code)
 
-    html_code = enhance_font_awesome_layout(html_code)
-    return html_code
 
 def trim_image_whitespace(image: Image.Image, threshold: int = 250, padding: int = 10) -> Image.Image:
-    """白余白トリミング（NumPy 高速化）"""
+    """白背景をトリミング"""
     gray = image.convert("L")
     arr = np.array(gray)
     mask = arr < threshold
-    rows = np.any(mask, axis=1)
-    cols = np.any(mask, axis=0)
+    rows, cols = np.any(mask, axis=1), np.any(mask, axis=0)
     if np.any(rows) and np.any(cols):
         y_min, y_max = np.where(rows)[0][[0, -1]]
         x_min, x_max = np.where(cols)[0][[0, -1]]
-        y_min = max(0, y_min - padding)
-        x_min = max(0, x_min - padding)
-        y_max = min(image.height - 1, y_max + padding)
-        x_max = min(image.width - 1,  x_max + padding)
-        return image.crop((x_min, y_min, x_max + 1, y_max + 1))
+        return image.crop((
+            max(0, x_min - padding), max(0, y_min - padding),
+            min(image.width - 1, x_max + padding), min(image.height - 1, y_max + padding)
+        ))
     return image
 
-# ----------------------------------------------------
+
+# ===============================================================
 # HTML → スクリーンショット
-# ----------------------------------------------------
+# ===============================================================
 def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0,
-                               trim_whitespace: bool = True,
-                               driver=None) -> Image.Image:
-    driver_from_pool = False
+                               trim_whitespace: bool = True, driver=None) -> Image.Image:
     tmp_path = None
+    driver_from_pool = False
     try:
         if driver is None:
             driver = driver_pool.get_driver()
             driver_from_pool = True
 
+        # 一時 HTML 保存
         with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode="w", encoding="utf-8") as tmp:
             tmp_path = tmp.name
             tmp.write(html_code)
 
-        initial_w, initial_h = 1200, 1000
-        driver.set_window_size(initial_w, initial_h)
-        driver.get(f"file://{tmp_path}")
+        driver.set_window_size(1200, 1000)
+        driver.get("file://" + tmp_path)
 
+        # body 待機
         WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.TAG_NAME, "body")))
 
-        total_h = driver.execute_script("return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)")
-        viewport_h = driver.execute_script("return window.innerHeight")
-        scrolls = max(1, min(5, total_h // viewport_h))
-        for i in range(scrolls):
-            driver.execute_script(f"window.scrollTo(0, {i * (viewport_h - 100)})")
+        total_height = driver.execute_script(
+            "return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)")
+        viewport_height = driver.execute_script("return window.innerHeight")
+        for i in range(max(1, min(5, total_height // viewport_height))):
+            driver.execute_script(f"window.scrollTo(0, {i * (viewport_height - 100)})")
             time.sleep(0.1)
-        driver.execute_script("window.scrollTo(0, 0)")
+        driver.execute_script("window.scrollTo(0,0)")
         time.sleep(0.2)
 
         dims = driver.execute_script("""
@@ -283,17 +289,14 @@ def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0
                 height: Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)
             }
         """)
-        w = min(max(dims["width"], 100), 2000)
-        h = min(max(dims["height"], 100), 4000)
-        h = int(h * (1 + extension_percentage / 100.0))
-        driver.set_window_size(w, h)
+        width = min(max(dims["width"], 100), 2000)
+        height = min(max(dims["height"], 100), 4000)
+        height = int(height * (1 + extension_percentage / 100.0))
+        driver.set_window_size(width, height)
         time.sleep(0.5)
 
-        png = driver.get_screenshot_as_png()
-        img = Image.open(BytesIO(png))
-        if trim_whitespace:
-            img = trim_image_whitespace(img, threshold=248, padding=20)
-        return img
+        img = Image.open(BytesIO(driver.get_screenshot_as_png()))
+        return trim_image_whitespace(img, 248, 20) if trim_whitespace else img
     except Exception as e:
         logger.error(f"Screenshot error: {e}", exc_info=True)
         return Image.new("RGB", (1, 1), (0, 0, 0))
@@ -306,28 +309,32 @@ def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0
             except Exception:
                 pass
 
-# ----------------------------------------------------
+
+# ===============================================================
 # テキスト → スクリーンショット（並列）
-# ----------------------------------------------------
+# ===============================================================
 def text_to_screenshot_parallel(text: str, extension_percentage: float, temperature: float = 0.5,
                                 trim_whitespace: bool = True, style: str = "standard") -> Image.Image:
     start = time.time()
-    with ThreadPoolExecutor(max_workers=2) as exe:
-        html_fut = exe.submit(generate_html_from_text, text, temperature, style)
-        driver_fut = exe.submit(driver_pool.get_driver)
+    with ThreadPoolExecutor(max_workers=2) as ex:
+        html_future = ex.submit(generate_html_from_text, text, temperature, style)
+        driver_future = ex.submit(driver_pool.get_driver)
 
-        html_code = html_fut.result()
-        driver = driver_fut.result()
+        html_code = html_future.result()
+        driver = driver_future.result()
     img = render_fullpage_screenshot(html_code, extension_percentage, trim_whitespace, driver)
-    logger.info(f"並列処理 完了 {time.time() - start:.2f}s")
+    logger.info(f"並列生成完了: {time.time() - start:.2f}s")
     return img
 
-def text_to_screenshot(*args, **kwargs) -> Image.Image:
+
+def text_to_screenshot(*args, **kwargs):
+    """後方互換用エイリアス"""
     return text_to_screenshot_parallel(*args, **kwargs)
 
-# ====================================================
-# FastAPI セットアップ
-# ====================================================
+
+# ===============================================================
+# FastAPI
+# ===============================================================
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -335,75 +342,100 @@ app.add_middleware(
     allow_methods=["*"], allow_headers=["*"]
 )
 
-# ------------ API エンドポイント ---------------
+
 @app.post("/api/screenshot", response_class=StreamingResponse, tags=["Screenshot"])
 async def api_render_screenshot(req: ScreenshotRequest):
     img = render_fullpage_screenshot(req.html_code, req.extension_percentage, req.trim_whitespace)
-    buf = BytesIO()
-    img.save(buf, format="PNG"); buf.seek(0)
+    buf = BytesIO(); img.save(buf, "PNG"); buf.seek(0)
     return StreamingResponse(buf, media_type="image/png")
 
+
 @app.post("/api/text-to-screenshot", response_class=StreamingResponse, tags=["Screenshot", "Gemini"])
 async def api_text_to_screenshot(req: GeminiRequest):
-    img = text_to_screenshot_parallel(req.text, req.extension_percentage,
-                                      req.temperature, req.trim_whitespace, req.style)
-    buf = BytesIO()
-    img.save(buf, format="PNG"); buf.seek(0)
+    img = text_to_screenshot_parallel(
+        req.text, req.extension_percentage, req.temperature, req.trim_whitespace, req.style)
+    buf = BytesIO(); img.save(buf, "PNG"); buf.seek(0)
     return StreamingResponse(buf, media_type="image/png")
 
-# ====================================================
+
+# ===============================================================
 # Gradio UI
-# ====================================================
-def process_input(mode, input_text, ext_perc, temp, trim_ws, style):
+# ===============================================================
+def process_input(mode, text, ext, temp, trim, style):
     if mode == "HTML入力":
-        return render_fullpage_screenshot(input_text, ext_perc, trim_ws)
-    return text_to_screenshot_parallel(input_text, ext_perc, temp, trim_ws, style)
+        return render_fullpage_screenshot(text, ext, trim)
+    return text_to_screenshot_parallel(text, ext, temp, trim, style)
+
 
 with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)",
                theme=gr.themes.Origin()) as iface:
-    gr.Markdown("# HTMLビューア & テキスト→インフォグラフィック")
+    gr.Markdown("# HTMLビューア & テキスト→インフォグラフィック変換")
     with gr.Row():
         input_mode = gr.Radio(["HTML入力", "テキスト入力"], value="HTML入力", label="入力モード")
+
     input_text = gr.Textbox(lines=15, label="入力")
+
     with gr.Row():
-        style_dropdown = gr.Dropdown(["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"],
-                                     value="standard", label="デザインスタイル", visible=False)
+        style_dropdown = gr.Dropdown(
+            ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"],
+            value="standard", label="デザインスタイル", visible=False)
         with gr.Column(scale=2):
             ext_slider = gr.Slider(0, 30, value=10, step=1, label="上下高さ拡張率（%）")
             temp_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.1,
                                     label="生成時の温度", visible=False)
-    trim_ws_chk = gr.Checkbox(value=True, label="余白を自動トリミング")
+
+    trim_chk = gr.Checkbox(value=True, label="余白を自動トリミング")
+
     submit_btn = gr.Button("生成")
     out_img = gr.Image(type="pil", label="スクリーンショット")
 
-    def toggle_controls(mode):
+    def toggle(mode):
         is_text = mode == "テキスト入力"
         return [gr.update(visible=is_text), gr.update(visible=is_text)]
-    input_mode.change(toggle_controls, input_mode, [temp_slider, style_dropdown])
-    submit_btn.click(process_input,
-                     [input_mode, input_text, ext_slider, temp_slider, trim_ws_chk, style_dropdown],
-                     out_img)
+    input_mode.change(toggle, input_mode, [temp_slider, style_dropdown])
+
+    submit_btn.click(
+        process_input,
+        [input_mode, input_text, ext_slider, temp_slider, trim_chk, style_dropdown],
+        out_img)
 
-    gemini_model = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro")
+    model_name = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro")
     gr.Markdown(f"""
 **API**  
 - `/api/screenshot` – HTML → PNG  
 - `/api/text-to-screenshot` – テキスト → インフォグラフィック PNG  
 
 **設定**  
-- 使用モデル: `{gemini_model}`  
-- スタイル: standard / cute / resort / cool / dental / school / KOKUGO  
+- 使用モデル: `{model_name}`  
+- 対応スタイル: standard / cute / resort / cool / dental / school / KOKUGO  
 - WebDriver 最大数: {driver_pool.max_drivers}
 """)
 
-# FastAPI へマウント（SSR 無効）
-app = gr.mount_gradio_app(app, iface, path="/", ssr_mode=False)
 
-# ローカル実行用
+# ===============================================================
+# FastAPI へマウント & ルートリダイレクト
+# ===============================================================
+GRADIO_PATH = "/gradio"
+app = gr.mount_gradio_app(app, iface, path=GRADIO_PATH, ssr_mode=False)
+
+
+@app.get("/")
+def _root():
+    """ルートに来たら /gradio へリダイレクト"""
+    return RedirectResponse(GRADIO_PATH)
+
+
+# ===============================================================
+# ローカルデバッグ用
+# ===============================================================
 if __name__ == "__main__":
     import uvicorn
+    logger.info("Uvicorn 起動 (ローカル)")
     uvicorn.run(app, host="0.0.0.0", port=7860)
 
-# 終了時のクリーンアップ
+
+# ===============================================================
+# 終了時 WebDriver 後始末
+# ===============================================================
 import atexit
 atexit.register(driver_pool.close_all)