HTMLviewer_Dev

Paused

App Files Files Community

tomo2chin2 commited on May 7, 2025

Commit

35a8897

verified ·

1 Parent(s): 64edcc1

Update app.py

Browse files

Files changed (1) hide show

app.py +72 -170

app.py CHANGED Viewed

@@ -6,7 +6,7 @@
 #   それ以外は 5.x 対応フルロジックを一切カットせず
 # ===============================================================
-import os, time, tempfile, logging, threading, queue, zipfile
 from io import BytesIO
 from concurrent.futures import ThreadPoolExecutor
@@ -126,22 +126,34 @@ class ScreenshotRequest(BaseModel):
     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]
@@ -149,10 +161,7 @@ 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")
-    prompt_text = ""
     if os.path.exists(local):
         prompt_text = open(local, encoding="utf-8").read()
     else:
@@ -167,50 +176,6 @@ def load_system_instruction(style="standard") -> str:
     _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 = """
-    <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>
-    """
-    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'{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"]
@@ -232,9 +197,9 @@ def generate_html_from_text(text: str, temperature=0.5, style="standard") -> str
     if model_name == "gemini-2.5-flash-preview-04-17":
         logger.info("gemini-2.5-flash-preview-04-17 モデル検出: 思考モードをオフに設定")
         config.thinking_config = types.ThinkingConfig(thinking_budget=0)
-        # max_output_tokens を 50000 に拡張
         logger.info("gemini-2.5-flash-preview-04-17 モデル検出: max_output_tokens を 50000 に設定")
-        config.max_output_tokens = 50000
     # Generate content
     response = client.models.generate_content(
@@ -261,7 +226,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:
@@ -272,62 +237,46 @@ 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)
-        # 非同期でリソースロード待機とスクリプト実行を行う
-        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
@@ -342,44 +291,18 @@ def render_fullpage_screenshot(html_code: str, extension_percentage=6.0,
             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 を殺す)
@@ -394,7 +317,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)
@@ -403,40 +326,17 @@ 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(
         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ビューア & テキスト→インフォグラフィック変換")
@@ -448,8 +348,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=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("生成")
@@ -465,7 +365,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 = "（思考モード: オフ、最大トークン: 50000）"
     gr.Markdown(f"**API** `/api/screenshot`, `/api/text-to-screenshot` &nbsp;&nbsp; "
                 f"使用モデル: `{model_name}` {thinking_status}")
@@ -484,12 +384,14 @@ def _root(): return RedirectResponse(GRADIO_PATH + "/")
 @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("アプリケーション起動: 並列処理による最適化を適用")
 # ===============================================================
 # ローカルデバッグ

 #   それ以外は 5.x 対応フルロジックを一切カットせず
 # ===============================================================
+import os, time, tempfile, logging, threading, queue
 from io import BytesIO
 from concurrent.futures import ThreadPoolExecutor
     trim_whitespace: bool = True
     style: str = "standard"
 # ---------------------------------------------------------------
+# 補助関数（FontAwesome レイアウト / prompt 読み込み / Gemini 生成）
 # ---------------------------------------------------------------
+def enhance_font_awesome_layout(html_code: str) -> str:
+    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>
+    """
+    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'{fa_preload}{fa_css}</head>')
+    return f'<html><head>{fa_preload}{fa_css}</head>{html_code}</html>'
+# シンプルなプロンプトキャッシュを実装
 _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")
     if os.path.exists(local):
         prompt_text = open(local, encoding="utf-8").read()
     else:
     _prompt_cache[style] = prompt_text
     return prompt_text
 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"]
     if model_name == "gemini-2.5-flash-preview-04-17":
         logger.info("gemini-2.5-flash-preview-04-17 モデル検出: 思考モードをオフに設定")
         config.thinking_config = types.ThinkingConfig(thinking_budget=0)
+        # max_output_tokens を 50000 に拡張 (唯一追加した最適化)
         logger.info("gemini-2.5-flash-preview-04-17 モデル検出: max_output_tokens を 50000 に設定")
+        config.max_output_tokens = 10000
     # Generate content
     response = client.models.generate_content(
     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)
+        # 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
             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 を殺す)
     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_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ビューア & テキスト→インフォグラフィック変換")
             ["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("生成")
     model_name = os.getenv('GEMINI_MODEL', 'gemini-1.5-pro')
     thinking_status = ""
     if model_name == "gemini-2.5-flash-preview-04-17":
+        thinking_status = "（思考モード: オフ、最大トークン: 10000）"
     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():
+    # 初期化は最小限に
+    styles = ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"]
+    for style in styles:
+        load_system_instruction(style)
+    logger.info("システムプロンプトのキャッシュを準備完了")
 # ===============================================================
 # ローカルデバッグ