HTMLviewer_Dev

Paused

App Files Files Community

tomo2chin2 commited on May 7, 2025

Commit

64edcc1

verified ·

1 Parent(s): a0f17f9

Update app.py

Browse files

Files changed (1) hide show

app.py +175 -56

app.py CHANGED Viewed

@@ -6,7 +6,7 @@
 #   それ以外は 5.x 対応フルロジックを一切カットせず
 # ===============================================================
-import os, time, tempfile, logging, threading, queue
 from io import BytesIO
 from concurrent.futures import ThreadPoolExecutor
@@ -126,8 +126,69 @@ class ScreenshotRequest(BaseModel):
     trim_whitespace: bool = True
     style: str = "standard"
 # ---------------------------------------------------------------
-# 補助関数（FontAwesome レイアウト / prompt 読み込み / Gemini 生成）
 # ---------------------------------------------------------------
 def enhance_font_awesome_layout(html_code: str) -> str:
     fa_preload = """
@@ -150,20 +211,6 @@ def enhance_font_awesome_layout(html_code: str) -> str:
         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") -> 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:
-        f = hf_hub_download("tomo2chin2/GURAREKOstlyle", f"{style}/prompt.txt", repo_type="dataset")
-        return open(f, encoding="utf-8").read()
-    except Exception:
-        f = hf_hub_download("tomo2chin2/GURAREKOstlyle", "prompt.txt", repo_type="dataset")
-        return open(f, encoding="utf-8").read()
 def generate_html_from_text(text: str, temperature=0.5, style="standard") -> str:
     # Updated: Use the new Google Genai client API
     api_key = os.environ["GEMINI_API_KEY"]
@@ -214,7 +261,7 @@ def trim_image_whitespace(img: Image.Image, threshold=248, padding=20) -> Image.
     return img
 # ---------------------------------------------------------------
-# HTML → スクショ  (完全版ロジック)
 # ---------------------------------------------------------------
 def render_fullpage_screenshot(html_code: str, extension_percentage=6.0,
                                trim_whitespace=True, driver=None) -> Image.Image:
@@ -225,46 +272,62 @@ def render_fullpage_screenshot(html_code: str, extension_percentage=6.0,
             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 + 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
@@ -279,18 +342,44 @@ def render_fullpage_screenshot(html_code: str, extension_percentage=6.0,
             except Exception: pass
 # ---------------------------------------------------------------
-# テキスト → スクショ (並列 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  (★ redirect_slashes=False で自動 307 を殺す)
@@ -305,7 +394,7 @@ app.add_middleware(
     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)
@@ -314,17 +403,40 @@ async def api_render_screenshot(req: ScreenshotRequest):
 @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 (完全版 UI 定義)
 # ===============================================================
 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ビューア & テキスト→インフォグラフィック変換")
@@ -336,8 +448,8 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
             ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"],
             value="standard", label="デザインスタイル", visible=False)
         with gr.Column(scale=2):
-            ext = gr.Slider(0, 30, value=15, step=1, label="上下高さ拡張率（%）")
-            temp = gr.Slider(0.0, 1.0, value=1.0, step=0.1,
                              label="生成時の温度", visible=False)
     trim = gr.Checkbox(value=True, label="余白を自動トリミング")
     btn = gr.Button("生成")
@@ -353,7 +465,7 @@ with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr
     model_name = os.getenv('GEMINI_MODEL', 'gemini-1.5-pro')
     thinking_status = ""
     if model_name == "gemini-2.5-flash-preview-04-17":
-        thinking_status = "（思考モード: オフ）"
     gr.Markdown(f"**API** `/api/screenshot`, `/api/text-to-screenshot` &nbsp;&nbsp; "
                 f"使用モデル: `{model_name}` {thinking_status}")
@@ -372,6 +484,13 @@ def _root(): return RedirectResponse(GRADIO_PATH + "/")
 @app.get(GRADIO_PATH)
 def _no_slash(): return RedirectResponse(GRADIO_PATH + "/")
 # ===============================================================
 # ローカルデバッグ
 # ===============================================================

 #   それ以外は 5.x 対応フルロジックを一切カットせず
 # ===============================================================
+import os, time, tempfile, logging, threading, queue, zipfile
 from io import BytesIO
 from concurrent.futures import ThreadPoolExecutor
     trim_whitespace: bool = True
     style: str = "standard"
+# バッチ処理用の新しいモデル
+class BatchGeminiRequest(BaseModel):
+    texts: list[str]
+    extension_percentage: float = 10.0
+    temperature: float = 0.5
+    trim_whitespace: bool = True
+    style: str = "standard"
+# ---------------------------------------------------------------
+# システム指示のキャッシュ実装
+# ---------------------------------------------------------------
+# プロンプトキャッシュ - 頻繁に使用されるプロンプトを保存
+_prompt_cache = {}
+def load_system_instruction(style="standard") -> str:
+    """システム指示をロード (キャッシュ機能付き)"""
+    # キャッシュに存在すればそれを返す
+    if style in _prompt_cache:
+        return _prompt_cache[style]
+    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")
+    prompt_text = ""
+    if os.path.exists(local):
+        prompt_text = open(local, encoding="utf-8").read()
+    else:
+        try:
+            f = hf_hub_download("tomo2chin2/GURAREKOstlyle", f"{style}/prompt.txt", repo_type="dataset")
+            prompt_text = open(f, encoding="utf-8").read()
+        except Exception:
+            f = hf_hub_download("tomo2chin2/GURAREKOstlyle", "prompt.txt", repo_type="dataset")
+            prompt_text = open(f, encoding="utf-8").read()
+    # キャッシュに保存
+    _prompt_cache[style] = prompt_text
+    return prompt_text
+# ---------------------------------------------------------------
+# 初期化時に全スタイルをキャッシュに先読み
+# ---------------------------------------------------------------
+def preload_all_prompts():
+    """アプリ起動時に全スタイルの指示を事前ロード"""
+    styles = ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"]
+    logger.info("システム指示の先読み開始...")
+    with ThreadPoolExecutor(max_workers=len(styles)) as executor:
+        futures = {executor.submit(load_system_instruction, style): style for style in styles}
+        for future in futures:
+            style = futures[future]
+            try:
+                future.result()  # 結果を取得
+                logger.info(f"スタイル '{style}' の指示を先読み完了")
+            except Exception as e:
+                logger.error(f"スタイル '{style}' の指示先読みに失敗: {e}")
+    logger.info(f"全 {len(_prompt_cache)} スタイルの指示先読み完了")
 # ---------------------------------------------------------------
+# 補助関数（FontAwesome レイアウト / Gemini 生成）
 # ---------------------------------------------------------------
 def enhance_font_awesome_layout(html_code: str) -> str:
     fa_preload = """
         return html_code.replace('</head>', f'{fa_preload}{fa_css}</head>')
     return f'<html><head>{fa_preload}{fa_css}</head>{html_code}</html>'
 def generate_html_from_text(text: str, temperature=0.5, style="standard") -> str:
     # Updated: Use the new Google Genai client API
     api_key = os.environ["GEMINI_API_KEY"]
     return img
 # ---------------------------------------------------------------
+# HTML → スクショ 最適化版 (並列処理強化)
 # ---------------------------------------------------------------
 def render_fullpage_screenshot(html_code: str, extension_percentage=6.0,
                                trim_whitespace=True, driver=None) -> Image.Image:
             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)
+        # 非同期でリソースロード待機とスクリプト実行を行う
+        def wait_for_resources():
+            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
+            return True
+        # リソース待機をスレッドプールで実行
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            resource_wait = executor.submit(wait_for_resources)
+            resource_wait.result()  # 待機完了を確認
+        # スクロールレンダリングを最適化
         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)
+        # サイズ計算と画像取得を並列化
+        def get_dimensions_and_resize():
+            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)
+            return w, h
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            dims_future = executor.submit(get_dimensions_and_resize)
+            dims_future.result()  # サイズ調整完了を確認
+        # スクリーンショット取得と画像処理
         img = Image.open(BytesIO(driver.get_screenshot_as_png()))
         return trim_image_whitespace(img, padding=20) if trim_whitespace else img
             except Exception: pass
 # ---------------------------------------------------------------
+# テキスト → スクショ (並列処理強化版)
 # ---------------------------------------------------------------
+def text_to_screenshot(text, ext_perc, temp=0.5, trim_ws=True, style="standard") -> Image.Image:
+    # 3つの並列タスク: HTML生成、ドライバ取得、必要なスタイルのシステム指示ロード
+    with ThreadPoolExecutor(max_workers=3) as exe:
+        # システム指示が未キャッシュの場合に備えて並列ロード
+        prompt_future = exe.submit(load_system_instruction, style)
         html_future = exe.submit(generate_html_from_text, text, temp, style)
         driver_future = exe.submit(driver_pool.get_driver)
+        # 結果取得
+        prompt_future.result()  # プロンプトをキャッシュ確保
         html_code = html_future.result()
         driver = driver_future.result()
+    # 最適化されたスクリーンショット関数を使用
     return render_fullpage_screenshot(html_code, ext_perc, trim_ws, driver)
+# ---------------------------------------------------------------
+# テキスト → スクショ (複数同時処理版)
+# ---------------------------------------------------------------
+def batch_text_to_screenshot(texts, ext_perc, temp=0.5, trim_ws=True, style="standard") -> list:
+    """複数テキストを同時に処理"""
+    with ThreadPoolExecutor(max_workers=min(len(texts), 3)) as exe:
+        futures = [exe.submit(text_to_screenshot, text, ext_perc, temp, trim_ws, style)
+                  for text in texts]
+        return [f.result() for f in futures]
+# ---------------------------------------------------------------
+# アプリ初期化時に実行する処理
+# ---------------------------------------------------------------
+def initialize_app():
+    """アプリケーション初期化処理"""
+    # システム指示を事前にキャッシュにロード
+    preload_all_prompts()
+    # その他の初期化処理があればここに追加
+    logger.info("アプリケーション初期化完了")
 # ===============================================================
 # FastAPI  (★ redirect_slashes=False で自動 307 を殺す)
     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)
 @app.post("/api/text-to-screenshot", response_class=StreamingResponse, tags=["Gemini","Screenshot"])
 async def api_text_to_screenshot(req: GeminiRequest):
+    img = text_to_screenshot(
         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")
+# バッチ処理用の新しいエンドポイント
+@app.post("/api/batch-text-to-screenshot", tags=["Gemini","Screenshot"])
+async def api_batch_text_to_screenshot(req: BatchGeminiRequest):
+    # 複数テキストを並列処理
+    images = batch_text_to_screenshot(
+        req.texts, req.extension_percentage, req.temperature, req.trim_whitespace, req.style)
+    # 結果をZIP形式で返す
+    buf = BytesIO()
+    with zipfile.ZipFile(buf, 'w') as zf:
+        for i, img in enumerate(images):
+            img_buf = BytesIO()
+            img.save(img_buf, format="PNG")
+            img_buf.seek(0)
+            zf.writestr(f"screenshot_{i+1}.png", img_buf.getvalue())
+    buf.seek(0)
+    return StreamingResponse(
+        buf,
+        media_type="application/zip",
+        headers={"Content-Disposition": "attachment; filename=screenshots.zip"}
+    )
 # ===============================================================
 # Gradio UI (完全版 UI 定義)
 # ===============================================================
 def process_input(mode, text, ext, temp, trim, style):
     return render_fullpage_screenshot(text, ext, trim) if mode == "HTML入力" else \
+           text_to_screenshot(text, ext, temp, trim, style)
 with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr.themes.Origin()) as iface:
     gr.Markdown("# HTMLビューア & テキスト→インフォグラフィック変換")
             ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"],
             value="standard", label="デザインスタイル", visible=False)
         with gr.Column(scale=2):
+            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("生成")
     model_name = os.getenv('GEMINI_MODEL', 'gemini-1.5-pro')
     thinking_status = ""
     if model_name == "gemini-2.5-flash-preview-04-17":
+        thinking_status = "（思考モード: オフ、最大トークン: 50000）"
     gr.Markdown(f"**API** `/api/screenshot`, `/api/text-to-screenshot` &nbsp;&nbsp; "
                 f"使用モデル: `{model_name}` {thinking_status}")
 @app.get(GRADIO_PATH)
 def _no_slash(): return RedirectResponse(GRADIO_PATH + "/")
+# アプリケーション起動時の初期化
+@app.on_event("startup")
+async def startup_event():
+    # バックグラウンドで初期化処理を実行
+    threading.Thread(target=initialize_app).start()
+    logger.info("アプリケーション起動: 並列処理による最適化を適用")
 # ===============================================================
 # ローカルデバッグ
 # ===============================================================