Update app.py

app.py

@@ -1,7 +1,9 @@
 # ===============================================================
-# app.py — Gradio 5.x + FastAPI + Gemini + Selenium
-#   • FastAPI(redirect_slashes=False) で 307 ループを根絶
-#   • サブアプリを /gradio にマウントし、自前でリダイレクト制御
+# app.py  ― Gradio 5.x + FastAPI + Gemini + Selenium + 307対策
+#   (1) FastAPI(redirect_slashes=False)               ←★追加
+#   (2) Gradio を /gradio にマウント
+#   (3) / と /gradio を /gradio/ へリダイレクト     ←★追加
+#   それ以外は 5.x 対応フルロジックを一切カットせず
 # ===============================================================
 
 import os, time, tempfile, logging, threading, queue
@@ -26,12 +28,17 @@ from selenium.webdriver.support import expected_conditions as EC
 import google.generativeai as genai
 from huggingface_hub import hf_hub_download
 
-# ---------- ログ ----------
+# ---------------------------------------------------------------
+# ロギング
+# ---------------------------------------------------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ---------- WebDriver プール（前回と同じ） ----------
+# ---------------------------------------------------------------
+# WebDriverPool  ― 以前提示した完全版と同一
+# ---------------------------------------------------------------
 class WebDriverPool:
+    """再利用可能な WebDriver をプールして高速化"""
     def __init__(self, max_drivers: int = 3):
         self.driver_queue = queue.Queue()
         self.max_drivers = max_drivers
@@ -40,40 +47,70 @@ class WebDriverPool:
         logger.info(f"WebDriver プール初期化: 最大 {max_drivers}")
 
     def _create_driver(self):
-        opts = Options()
-        opts.add_argument("--headless"); opts.add_argument("--no-sandbox")
-        opts.add_argument("--disable-dev-shm-usage")
-        drv_path = os.getenv("CHROMEDRIVER_PATH")
-        if drv_path and os.path.exists(drv_path):
-            return webdriver.Chrome(service=webdriver.ChromeService(executable_path=drv_path), options=opts)
-        return webdriver.Chrome(options=opts)
+        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):
         if not self.driver_queue.empty():
+            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})")
                 return self._create_driver()
+
+        logger.info("プール満杯、空き待機中…")
         return self.driver_queue.get()
 
     def release_driver(self, driver):
-        try:
-            driver.get("about:blank")
-            self.driver_queue.put(driver)
-        except Exception:
-            driver.quit()
-            with self.lock: self.count -= 1
+        if driver:
+            try:
+                driver.get("about:blank")
+                driver.execute_script("""
+                    document.documentElement.style.overflow='';
+                    document.body.style.overflow='';
+                """)
+                self.driver_queue.put(driver)
+                logger.info("WebDriver をプールに返却")
+            except Exception as e:
+                logger.error(f"返却エラー: {e}")
+                driver.quit()
+                with self.lock:
+                    self.count -= 1
 
     def close_all(self):
+        logger.info("プール全 WebDriver 終了")
+        closed = 0
         while not self.driver_queue.empty():
-            try: self.driver_queue.get(block=False).quit()
-            except queue.Empty: break
+            try:
+                drv = self.driver_queue.get(block=False)
+                drv.quit(); closed += 1
+            except queue.Empty:
+                break
+            except Exception as e:
+                logger.error(f"終了エラー: {e}")
         with self.lock: self.count = 0
+        logger.info(f"{closed} 個終了")
 
 driver_pool = WebDriverPool(max_drivers=int(os.getenv("MAX_WEBDRIVERS", "3")))
 
-# ---------- Pydantic モデル ----------
+# ---------------------------------------------------------------
+# Pydantic モデル
+# ---------------------------------------------------------------
 class GeminiRequest(BaseModel):
     text: str
     extension_percentage: float = 10.0
@@ -87,147 +124,229 @@ class ScreenshotRequest(BaseModel):
     trim_whitespace: bool = True
     style: str = "standard"
 
-# ---------- 補助関数（前回と同じ：FontAwesome, prompt 取得, Gemini, トリミング） ----------
+# ---------------------------------------------------------------
+# 補助関数（FontAwesome レイアウト / prompt 読み込み / Gemini 生成）
+# ---------------------------------------------------------------
 def enhance_font_awesome_layout(html_code: str) -> str:
-    preload = """
+    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 = '<style>[class*="fa-"]{display:inline-block;margin-right:8px;vertical-align:middle}</style>'
+    fa_css = """
+    <style>
+      [class*="fa-"]{display:inline-block!important;margin-right:8px!important;vertical-align:middle!important;}
+      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>
+    """
     if '<head>' in html_code:
-        return html_code.replace('</head>', f'{preload}{css}</head>')
-    return f'<html><head>{preload}{css}</head>{html_code}</html>'
+        return html_code.replace('</head>', f'{fa_preload}{fa_css}</head>')
+    return f'<html><head>{fa_preload}{fa_css}</head>{html_code}</html>'
 
-def load_system_instruction(style="standard"):
-    if style not in ["standard","cute","resort","cool","dental","school","KOKUGO"]:
+def load_system_instruction(style="standard") -> str:
+    valid_styles = ["standard","cute","resort","cool","dental","school","KOKUGO"]
+    if style not in valid_styles:
         style = "standard"
     local = os.path.join(os.path.dirname(__file__), style, "prompt.txt")
     if os.path.exists(local):
         return open(local, encoding="utf-8").read()
     try:
-        path = hf_hub_download("tomo2chin2/GURAREKOstlyle", f"{style}/prompt.txt", repo_type="dataset")
-        return open(path, encoding="utf-8").read()
+        f = hf_hub_download("tomo2chin2/GURAREKOstlyle", f"{style}/prompt.txt", repo_type="dataset")
+        return open(f, encoding="utf-8").read()
     except Exception:
-        path = hf_hub_download("tomo2chin2/GURAREKOstlyle", "prompt.txt", repo_type="dataset")
-        return open(path, encoding="utf-8").read()
+        f = hf_hub_download("tomo2chin2/GURAREKOstlyle", "prompt.txt", repo_type="dataset")
+        return open(f, encoding="utf-8").read()
 
-def generate_html_from_text(text, temperature=0.5, style="standard"):
-    genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
-    model = genai.GenerativeModel(os.getenv("GEMINI_MODEL","gemini-1.5-pro"))
+def generate_html_from_text(text: str, temperature=0.5, style="standard") -> str:
+    api_key = os.environ["GEMINI_API_KEY"]
+    genai.configure(api_key=api_key)
+    model = genai.GenerativeModel(os.getenv("GEMINI_MODEL", "gemini-1.5-pro"))
     prompt = f"{load_system_instruction(style)}\n\n{text}"
-    cfg=dict(temperature=temperature,top_p=0.7,top_k=20,max_output_tokens=8192,candidate_count=1)
-    raw=model.generate_content(prompt,generation_config=cfg).text
-    s,e=raw.find("```html"),raw.rfind("```")
-    html=raw[s+7:e].strip() if s!=-1 and e!=-1 else raw
+    cfg = dict(temperature=temperature, top_p=0.7, top_k=20, max_output_tokens=8192, candidate_count=1)
+    resp = model.generate_content(prompt, generation_config=cfg)
+    raw = resp.text
+    s, e = raw.find("```html"), raw.rfind("```")
+    html = raw[s+7:e].strip() if s != -1 and e != -1 else raw
     return enhance_font_awesome_layout(html)
 
-def trim_image_whitespace(img:Image.Image,th=248,pad=20)->Image.Image:
-    arr=np.array(img.convert("L")); m=arr<th
-    if m.any():
-        y,x=np.where(m.any(1))[0],np.where(m.any(0))[0]
-        return img.crop((max(0,x[0]-pad),max(0,y[0]-pad),
-                         min(img.width-1,x[-1]+pad),min(img.height-1,y[-1]+pad)))
+def trim_image_whitespace(img: Image.Image, threshold=248, padding=20) -> Image.Image:
+    arr = np.array(img.convert("L"))
+    mask = arr < threshold
+    if np.any(mask):
+        ys, xs = np.where(mask.any(1))[0], np.where(mask.any(0))[0]
+        return img.crop((max(xs[0]-padding,0), max(ys[0]-padding,0),
+                         min(xs[-1]+padding, img.width-1),
+                         min(ys[-1]+padding, img.height-1)))
     return img
 
-# ---------- HTML → PNG ----------
-def render_fullpage_screenshot(html, ext=6.0, trim_ws=True, driver=None):
-    tmp=None; from_pool=False
+# ---------------------------------------------------------------
+# HTML → スクショ  (完全版ロジック)
+# ---------------------------------------------------------------
+def render_fullpage_screenshot(html_code: str, extension_percentage=6.0,
+                               trim_whitespace=True, driver=None) -> Image.Image:
+    tmp_path = None
+    from_pool = False
     try:
         if driver is None:
-            driver=driver_pool.get_driver(); from_pool=True
-        with tempfile.NamedTemporaryFile(suffix=".html",delete=False,mode="w",encoding="utf-8") as f:
-            tmp=f.name; f.write(html)
-        driver.set_window_size(1200,1000)
-        driver.get("file://"+tmp)
-        WebDriverWait(driver,10).until(EC.presence_of_element_located((By.TAG_NAME,"body")))
-        h=driver.execute_script("return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)")
-        vh=driver.execute_script("return window.innerHeight")
-        for i in range(max(1,min(5,h//vh))):
-            driver.execute_script(f"window.scrollTo(0,{i*(vh-100)})"); time.sleep(0.1)
+            driver = driver_pool.get_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)
+
+        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")))
+
+        # リソースロード確認ループ（詳細ロジックは元コード準拠）
+        max_wait, inc, waited = 5, 0.2, 0.0
+        while waited < max_wait:
+            state = driver.execute_script("""
+                return {complete: document.readyState==='complete',
+                        imgs: document.images.length,
+                        loaded: Array.from(document.images).filter(i=>i.complete).length};
+            """)
+            if state['complete'] and (state['imgs']==0 or state['imgs']==state['loaded']):
+                break
+            time.sleep(inc); waited += inc
+
+        # スクロールレンダリング
+        total_h = driver.execute_script("return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)")
+        vh = driver.execute_script("return window.innerHeight")
+        for i in range(max(1, min(5, total_h // vh))):
+            driver.execute_script(f"window.scrollTo(0, {(vh-100)*i})")
+            time.sleep(0.1)
         driver.execute_script("window.scrollTo(0,0)"); time.sleep(0.2)
-        dims=driver.execute_script("return {w:Math.max(document.body.scrollWidth,document.documentElement.scrollWidth),h:Math.max(document.body.scrollHeight,document.documentElement.scrollHeight)}")
-        w=min(max(dims['w'],100),2000); h=min(max(dims['h'],100),4000); h=int(h*(1+ext/100))
-        driver.set_window_size(w,h); time.sleep(0.4)
-        img=Image.open(BytesIO(driver.get_screenshot_as_png()))
-        return trim_image_whitespace(img,248,20) if trim_ws else img
+
+        dims = driver.execute_script("""
+            return {w: Math.max(document.body.scrollWidth, document.documentElement.scrollWidth),
+                    h: Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)}
+        """)
+        w = min(max(dims['w'], 100), 2000)
+        h = min(max(dims['h'], 100), 4000)
+        h = int(h * (1 + extension_percentage / 100.0))
+        driver.set_window_size(w, h); time.sleep(0.5)
+
+        img = Image.open(BytesIO(driver.get_screenshot_as_png()))
+        return trim_image_whitespace(img, padding=20) if trim_whitespace else img
+
     except Exception as e:
-        logger.error(f"Screenshot err: {e}",exc_info=True)
-        return Image.new("RGB",(1,1),(0,0,0))
+        logger.error(f"Screenshot error: {e}", exc_info=True)
+        return Image.new("RGB", (1, 1), (0, 0, 0))
     finally:
-        if from_pool: driver_pool.release_driver(driver)
-        if tmp and os.path.exists(tmp): os.remove(tmp)
+        if from_pool:
+            driver_pool.release_driver(driver)
+        if tmp_path and os.path.exists(tmp_path):
+            try: os.remove(tmp_path)
+            except Exception: pass
 
-# ---------- テキスト → PNG ----------
-def text_to_screenshot_parallel(text, ext, temp=0.5, trim_ws=True, style="standard"):
-    with ThreadPoolExecutor(max_workers=2) as ex:
-        html=ex.submit(generate_html_from_text,text,temp,style).result()
-        drv=ex.submit(driver_pool.get_driver).result()
-    return render_fullpage_screenshot(html,ext,trim_ws,drv)
+# ---------------------------------------------------------------
+# テキスト → スクショ (並列 API 呼び出し + ドライバ確保)
+# ---------------------------------------------------------------
+def text_to_screenshot_parallel(text, ext_perc, temp=0.5, trim_ws=True, style="standard") -> Image.Image:
+    with ThreadPoolExecutor(max_workers=2) as exe:
+        html_future = exe.submit(generate_html_from_text, text, temp, style)
+        driver_future = exe.submit(driver_pool.get_driver)
+        html_code = html_future.result()
+        driver = driver_future.result()
+    return render_fullpage_screenshot(html_code, ext_perc, trim_ws, driver)
+
+def text_to_screenshot(*args, **kwargs):
+    return text_to_screenshot_parallel(*args, **kwargs)
 
 # ===============================================================
-# FastAPI （★ リダイレクト無効化）
+# FastAPI  (★ redirect_slashes=False で自動 307 を殺す)
 # ===============================================================
 app = FastAPI(redirect_slashes=False)
-app.add_middleware(CORSMiddleware,allow_origins=["*"],allow_credentials=True,
-                   allow_methods=["*"],allow_headers=["*"])
-
-@app.post("/api/screenshot",response_class=StreamingResponse,tags=["Screenshot"])
-async def api_screenshot(r:ScreenshotRequest):
-    img=render_fullpage_screenshot(r.html_code,r.extension_percentage,r.trim_whitespace)
-    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=["Gemini","Screenshot"])
-async def api_text_to_screenshot(r:GeminiRequest):
-    img=text_to_screenshot_parallel(r.text,r.extension_percentage,r.temperature,r.trim_whitespace,r.style)
-    buf=BytesIO(); img.save(buf,"PNG"); buf.seek(0); return StreamingResponse(buf,media_type="image/png")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    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)
+    return StreamingResponse(buf, media_type="image/png")
+
+@app.post("/api/text-to-screenshot", response_class=StreamingResponse, tags=["Gemini","Screenshot"])
+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)
+    return StreamingResponse(buf, media_type="image/png")
 
 # ===============================================================
-# Gradio UI
+# Gradio UI (完全版 UI 定義)
 # ===============================================================
-def process(mode,txt,ext,temp,trim,style):
-    return render_fullpage_screenshot(txt,ext,trim) if mode=="HTML入力" else \
-           text_to_screenshot_parallel(txt,ext,temp,trim,style)
-
-with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)",
-               theme=gr.themes.Origin()) as gui:
-    gr.Markdown("# HTMLビューア & テキスト→インフォグラフィック")
-    mode=gr.Radio(["HTML入力","テキスト入力"],value="HTML入力",label="入力モード")
-    txt=gr.Textbox(lines=15,label="入力")
+def process_input(mode, text, ext, temp, trim, style):
+    return render_fullpage_screenshot(text, ext, trim) if mode == "HTML入力" else \
+           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ビューア & テキスト→インフォグラフィック変換")
+    with gr.Row():
+        mode = gr.Radio(["HTML入力", "テキスト入力"], value="HTML入力", label="入力モード")
+    text = gr.Textbox(lines=15, label="入力")
     with gr.Row():
-        style=gr.Dropdown(["standard","cute","resort","cool","dental","school","KOKUGO"],
-                          value="standard",label="デザインスタイル",visible=False)
+        style = gr.Dropdown(
+            ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"],
+            value="standard", label="デザインスタイル", visible=False)
         with gr.Column(scale=2):
-            ext=gr.Slider(0,30,10,1,label="上下高さ拡張率（%）")
-            temp=gr.Slider(0.0,1.0,0.5,0.1,label="生成時の温度",visible=False)
-    trim=gr.Checkbox(True,label="余白を自動トリミング")
-    btn=gr.Button("生成")
-    out=gr.Image(type="pil",label="スクリーンショット")
+            ext = gr.Slider(0, 30, value=10, step=1, label="上下高さ拡張率（%）")
+            temp = gr.Slider(0.0, 1.0, value=0.5, step=0.1,
+                             label="生成時の温度", visible=False)
+    trim = gr.Checkbox(value=True, label="余白を自動トリミング")
+    btn = gr.Button("生成")
+    out = gr.Image(type="pil", label="スクリーンショット")
 
-    mode.change(lambda m:[gr.update(visible=m=="テキスト入力")]*2,mode,[temp,style])
-    btn.click(process,[mode,txt,ext,temp,trim,style],out)
+    # 可視制御
+    def toggle(m): vis = m == "テキスト入力"; return [gr.update(visible=vis), gr.update(visible=vis)]
+    mode.change(toggle, mode, [temp, style])
 
-    gr.Markdown(f"**API**  `/api/screenshot`, `/api/text-to-screenshot`  \n使用モデル: `{os.getenv('GEMINI_MODEL','gemini-1.5-pro')}`")
+    btn.click(process_input, [mode, text, ext, temp, trim, style], out)
 
-# ---------------------------------------------------------------
-# Gradio を /gradio にマウント
-# ---------------------------------------------------------------
-GRADIO_BASE="/gradio"
-app = gr.mount_gradio_app(app, gui, path=GRADIO_BASE, ssr_mode=False)
+    gr.Markdown(f"**API** `/api/screenshot`, `/api/text-to-screenshot` &nbsp;&nbsp; "
+                f"使用モデル: `{os.getenv('GEMINI_MODEL','gemini-1.5-pro')}`")
 
-# ルート → /gradio/ へワンショット転送
+# ===============================================================
+# Gradio を /gradio にマウントし、明示リダイレクトを追加
+# ===============================================================
+GRADIO_PATH = "/gradio"
+app = gr.mount_gradio_app(app, iface, path=GRADIO_PATH, ssr_mode=False)
+
+# ルート → /gradio/ へ転送
 @app.get("/")
-def _root(): return RedirectResponse(GRADIO_BASE+"/")
+def _root(): return RedirectResponse(GRADIO_PATH + "/")
 
-# /gradio → /gradio/ にも転送して 307 を回避
-@app.get(GRADIO_BASE)
-def _no_slash(): return RedirectResponse(GRADIO_BASE+"/")
+# 末尾スラッシュ無し → 有りへ転送
+@app.get(GRADIO_PATH)
+def _no_slash(): return RedirectResponse(GRADIO_PATH + "/")
 
+# ===============================================================
 # ローカルデバッグ
-if __name__=="__main__":
-    import uvicorn, httpx
+# ===============================================================
+if __name__ == "__main__":
+    import uvicorn
     logger.info("Uvicorn 起動 (ローカル)")
-    uvicorn.run(app,host="0.0.0.0",port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)
 
-# 終了クリーンアップ
-import atexit; atexit.register(driver_pool.close_all)
+# ===============================================================
+# 終了時 WebDriver クリーンアップ
+# ===============================================================
+import atexit
+atexit.register(driver_pool.close_all)