# =============================================================== # 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 from io import BytesIO from concurrent.futures import ThreadPoolExecutor import numpy as np from PIL import Image import gradio as gr from fastapi import FastAPI, HTTPException from fastapi.responses import StreamingResponse, RedirectResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel 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 # Updated: Use the new Google Genai library from google import genai from google.genai import types from huggingface_hub import hf_hub_download # --------------------------------------------------------------- # ロギング # --------------------------------------------------------------- logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # --------------------------------------------------------------- # WebDriverPool ― 以前提示した完全版と同一 # --------------------------------------------------------------- class WebDriverPool: """再利用可能な WebDriver をプールして高速化""" def __init__(self, max_drivers: int = 3): self.driver_queue = queue.Queue() self.max_drivers = max_drivers self.lock = threading.Lock() self.count = 0 logger.info(f"WebDriver プール初期化: 最大 {max_drivers}") def _create_driver(self): 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): 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: 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 モデル # --------------------------------------------------------------- class GeminiRequest(BaseModel): text: str extension_percentage: float = 10.0 temperature: float = 0.5 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" # --------------------------------------------------------------- # 補助関数（FontAwesome レイアウト / prompt 読み込み / Gemini 生成） # --------------------------------------------------------------- def enhance_font_awesome_layout(html_code: str) -> str: fa_preload = """ """ fa_css = """ """ if '' in html_code: return html_code.replace('', f'{fa_preload}{fa_css}') return f'{fa_preload}{fa_css}{html_code}' # シンプルなプロンプトキャッシュを実装 _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: 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 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"] client = genai.Client(api_key=api_key) model_name = os.getenv("GEMINI_MODEL", "gemini-1.5-pro") prompt = f"{load_system_instruction(style)}\n\n{text}" # Configure generation parameters config = types.GenerateContentConfig( temperature=temperature, top_p=0.7, top_k=20, max_output_tokens=8192, candidate_count=1 ) # Gemini 2.5 Flash Preview モデルの場合の特別設定 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( model=model_name, contents=prompt, config=config ) # Extract HTML from response raw = response.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, 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 → スクショ (完全版ロジック) # --------------------------------------------------------------- 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 # 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 as e: 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_path and os.path.exists(tmp_path): try: os.remove(tmp_path) 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 を殺す) # =============================================================== app = FastAPI(redirect_slashes=False) 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 (完全版 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ビューア & テキスト→インフォグラフィック変換") 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) 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("生成") out = gr.Image(type="pil", label="スクリーンショット") # 可視制御 def toggle(m): vis = m == "テキスト入力"; return [gr.update(visible=vis), gr.update(visible=vis)] mode.change(toggle, mode, [temp, style]) btn.click(process_input, [mode, text, ext, temp, trim, style], out) # モデル名を表示し、Gemini 2.5の場合は思考モードのステータスも表示 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`    " f"使用モデル: `{model_name}` {thinking_status}") # =============================================================== # 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_PATH + "/") # 末尾スラッシュ無し → 有りへ転送 @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("システムプロンプトのキャッシュを準備完了") # =============================================================== # ローカルデバッグ # =============================================================== 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)