import gradio as gr from fastapi import FastAPI, HTTPException, Body from fastapi.responses import StreamingResponse from fastapi.staticfiles import StaticFiles 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 from PIL import Image from io import BytesIO import tempfile import time import os import logging import numpy as np # 追加: 画像処理の最適化用 import threading # 追加: 並列処理のため import queue # 追加: WebDriverプール用 from concurrent.futures import ThreadPoolExecutor # 追加: 並列処理用 from huggingface_hub import hf_hub_download # 正しいGemini関連のインポート import google.generativeai as genai # ロギング設定 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # --- WebDriverプールの実装 --- class WebDriverPool: """WebDriverインスタンスを再利用するためのプール""" def __init__(self, max_drivers=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 get_driver(self): """プールからWebDriverを取得、なければ新規作成""" 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})") 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") # 環境変数からWebDriverパスを取得（任意） webdriver_path = os.environ.get("CHROMEDRIVER_PATH") if webdriver_path and os.path.exists(webdriver_path): logger.info(f"CHROMEDRIVER_PATH使用: {webdriver_path}") service = webdriver.ChromeService(executable_path=webdriver_path) return webdriver.Chrome(service=service, options=options) else: logger.info("デフォルトのChromeDriverを使用") return webdriver.Chrome(options=options) # 最大数に達した場合は待機 logger.info("WebDriverプールがいっぱいです。利用可能なドライバーを待機中...") 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: driver = self.driver_queue.get(block=False) driver.quit() closed += 1 except queue.Empty: break except Exception as e: logger.error(f"ドライバー終了中にエラー: {e}") logger.info(f"{closed}個のWebDriverを終了しました") with self.lock: self.count = 0 # グローバルなWebDriverプールを作成 # サーバー環境のリソースに合わせて調整 driver_pool = WebDriverPool(max_drivers=int(os.environ.get("MAX_WEBDRIVERS", "3"))) # --- Gemini統合 --- class GeminiRequest(BaseModel): """Geminiへのリクエストデータモデル""" text: str extension_percentage: float = 10.0 # デフォルト値10% temperature: float = 0.5 # デフォルト値を0.5に設定 trim_whitespace: bool = True # 余白トリミングオプション（デフォルト有効） style: str = "standard" # デフォルトはstandard class ScreenshotRequest(BaseModel): """スクリーンショットリクエストモデル""" html_code: str extension_percentage: float = 10.0 # デフォルト値10% trim_whitespace: bool = True # 余白トリミングオプション（デフォルト有効） style: str = "standard" # デフォルトはstandard # HTMLのFont Awesomeレイアウトを改善する関数 - プリロード機能を追加 def enhance_font_awesome_layout(html_code): """Font Awesomeレイアウトを改善し、プリロードタグを追加""" # Font Awesomeリソースのプリロード - パフォーマンス向上 fa_preload = """ """ # CSSを追加 fa_fix_css = """ """ # headタグがある場合はその中に追加 if '' in html_code: return html_code.replace('', f'{fa_preload}{fa_fix_css}') # HTMLタグがある場合はその後に追加 elif '') if head_end > 0: return html_code[:head_end] + fa_preload + fa_fix_css + html_code[head_end:] else: body_start = html_code.find(' 0: return html_code[:body_start] + f'{fa_preload}{fa_fix_css}' + html_code[body_start:] # どちらもない場合は先頭に追加 return f'{fa_preload}{fa_fix_css}' + html_code + '' def load_system_instruction(style="standard"): """ 指定されたスタイルのシステムインストラクションを読み込む Args: style: 使用するスタイル名 (standard, cute, resort, cool, dental, school) Returns: 読み込まれたシステムインストラクション """ try: # 有効なスタイル一覧 valid_styles = ["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"] # スタイルの検証 if style not in valid_styles: logger.warning(f"無効なスタイル '{style}' が指定されました。デフォルトの 'standard' を使用します。") style = "standard" logger.info(f"スタイル '{style}' のシステムインストラクションを読み込みます") # まず、ローカルのスタイルディレクトリ内のprompt.txtを確認 local_path = os.path.join(os.path.dirname(__file__), style, "prompt.txt") # ローカルファイルが存在する場合はそれを使用 if os.path.exists(local_path): logger.info(f"ローカルファイルを使用: {local_path}") with open(local_path, 'r', encoding='utf-8') as file: instruction = file.read() return instruction # HuggingFaceリポジトリからのファイル読み込みを試行 try: # スタイル固有のファイルパスを指定 file_path = hf_hub_download( repo_id="tomo2chin2/GURAREKOstlyle", filename=f"{style}/prompt.txt", repo_type="dataset" ) logger.info(f"スタイル '{style}' のプロンプトをHuggingFaceから読み込みました: {file_path}") with open(file_path, 'r', encoding='utf-8') as file: instruction = file.read() return instruction except Exception as style_error: # スタイル固有ファイルの読み込みに失敗した場合、デフォルトのprompt.txtを使用 logger.warning(f"スタイル '{style}' のプロンプト読み込みに失敗: {str(style_error)}") logger.info("デフォルトのprompt.txtを読み込みます") file_path = hf_hub_download( repo_id="tomo2chin2/GURAREKOstlyle", filename="prompt.txt", repo_type="dataset" ) with open(file_path, 'r', encoding='utf-8') as file: instruction = file.read() logger.info("デフォルトのシステムインストラクションを読み込みました") return instruction except Exception as e: error_msg = f"システムインストラクションの読み込みに失敗: {str(e)}" logger.error(error_msg) raise ValueError(error_msg) def generate_html_from_text(text, temperature=0.5, style="standard"): """テキストからHTMLを生成する""" try: # APIキーの取得と設定 api_key = os.environ.get("GEMINI_API_KEY") if not api_key: logger.error("GEMINI_API_KEY 環境変数が設定されていません") raise ValueError("GEMINI_API_KEY 環境変数が設定されていません") # モデル名の取得（環境変数から、なければデフォルト値） model_name = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro") logger.info(f"使用するGeminiモデル: {model_name}") # Gemini APIの設定 genai.configure(api_key=api_key) # 指定されたスタイルのシステムインストラクションを読み込む system_instruction = load_system_instruction(style) # モデル初期化 logger.info(f"Gemini APIにリクエストを送信: テキスト長さ = {len(text)}, 温度 = {temperature}, スタイル = {style}") # モデル初期化 model = genai.GenerativeModel(model_name) # 生成設定 - ばらつきを減らすために設定を調整 generation_config = { "temperature": temperature, # より低い温度を設定 "top_p": 0.7, # 0.95から0.7に下げて出力の多様性を制限 "top_k": 20, # 64から20に下げて候補を絞る "max_output_tokens": 8192, "candidate_count": 1 # 候補は1つだけ生成 } # 安全設定 - デフォルトの安全設定を使用 safety_settings = [ { "category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE" }, { "category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE" }, { "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 ) # レスポンスからHTMLを抽出 raw_response = response.text # HTMLタグ部分だけを抽出（```html と ``` の間） html_start = raw_response.find("```html") html_end = raw_response.rfind("```") if html_start != -1 and html_end != -1 and html_start < html_end: html_start += 7 # "```html" の長さ分進める html_code = raw_response[html_start:html_end].strip() logger.info(f"HTMLの生成に成功: 長さ = {len(html_code)}") # Font Awesomeのレイアウト改善 html_code = enhance_font_awesome_layout(html_code) logger.info("Font Awesomeレイアウトの最適化を適用しました") return html_code else: # HTMLタグが見つからない場合、レスポンス全体を返す logger.warning("レスポンスから ```html ``` タグが見つかりませんでした。全テキストを返します。") return raw_response except Exception as e: logger.error(f"HTML生成中にエラーが発生: {e}", exc_info=True) raise Exception(f"Gemini APIでのHTML生成に失敗しました: {e}") # 画像から余分な空白領域をトリミングする関数 - NumPyを使って最適化 def trim_image_whitespace(image, threshold=250, padding=10): """ NumPyを使用して最適化された画像トリミング関数 Args: image: PIL.Image - 入力画像 threshold: int - どの明るさ以上を空白と判断するか (0-255) padding: int - トリミング後に残す余白のピクセル数 Returns: トリミングされたPIL.Image """ try: # グレースケールに変換 gray = image.convert('L') # NumPy配列として取得（高速処理のため） np_image = np.array(gray) # マスク作成（非白ピクセル） mask = np_image < threshold # マスクから行と列のインデックスを取得 rows = np.any(mask, axis=1) cols = np.any(mask, axis=0) # 非空のインデックス範囲を取得 if np.any(rows) and np.any(cols): row_indices = np.where(rows)[0] col_indices = np.where(cols)[0] # 範囲取得 min_y, max_y = row_indices[0], row_indices[-1] min_x, max_x = col_indices[0], col_indices[-1] # パディング追加 min_x = max(0, min_x - padding) min_y = max(0, min_y - padding) max_x = min(image.width - 1, max_x + padding) max_y = min(image.height - 1, max_y + padding) # 画像をトリミング trimmed = image.crop((min_x, min_y, max_x + 1, max_y + 1)) logger.info(f"画像をトリミングしました: 元サイズ {image.width}x{image.height} → トリミング後 {trimmed.width}x{trimmed.height}") return trimmed logger.warning("トリミング領域が見つかりません。元の画像を返します。") return image except Exception as e: logger.error(f"画像トリミング中にエラー: {e}", exc_info=True) return image # エラー時は元の画像を返す # 最適化されたスクリーンショット生成関数 - 外部から初期化済みドライバーを受け取れるように def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0, trim_whitespace: bool = True, driver=None) -> Image.Image: """ Renders HTML code to a full-page screenshot using Selenium. Optimized to accept an external driver or get one from the pool. Args: html_code: The HTML source code string. extension_percentage: Percentage of extra space to add vertically. trim_whitespace: Whether to trim excess whitespace from the image. driver: An optional pre-initialized WebDriver instance. Returns: A PIL Image object of the screenshot. """ tmp_path = None driver_from_pool = False # ドライバーがない場合はプールから取得 if driver is None: driver = driver_pool.get_driver() driver_from_pool = True logger.info("WebDriverプールからドライバーを取得しました") # 1) Save HTML code to a temporary file try: with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode='w', encoding='utf-8') as tmp_file: tmp_path = tmp_file.name tmp_file.write(html_code) logger.info(f"HTML saved to temporary file: {tmp_path}") except Exception as e: logger.error(f"Error writing temporary HTML file: {e}") if driver_from_pool: driver_pool.release_driver(driver) return Image.new('RGB', (1, 1), color=(0, 0, 0)) try: # ウィンドウサイズ初期設定 initial_width = 1200 initial_height = 1000 driver.set_window_size(initial_width, initial_height) file_url = "file://" + tmp_path logger.info(f"Navigating to {file_url}") driver.get(file_url) # ページ読み込み待機 - 動的な待機時間を実装 logger.info("Waiting for body element...") WebDriverWait(driver, 10).until( EC.presence_of_element_located((By.TAG_NAME, "body")) ) logger.info("Body element found. Waiting for resource loading...") # リソース読み込みの動的待機 - 最適化 max_wait = 5 # 最大待機時間（秒） wait_increment = 0.2 # 確認間隔 wait_time = 0 while wait_time < max_wait: resource_state = driver.execute_script(""" return { complete: document.readyState === 'complete', imgCount: document.images.length, imgLoaded: Array.from(document.images).filter(img => img.complete).length, faElements: document.querySelectorAll('.fa, .fas, .far, .fab, [class*="fa-"]').length }; """) # ドキュメント完了かつ画像が読み込まれている場合、待機終了 if resource_state['complete'] and (resource_state['imgCount'] == 0 or resource_state['imgLoaded'] == resource_state['imgCount']): logger.info(f"リソース読み込み完了: {resource_state}") break time.sleep(wait_increment) wait_time += wait_increment logger.info(f"リソース待機中... {wait_time:.1f}秒経過, 状態: {resource_state}") # Font Awesome要素が多い場合は追加待機 fa_count = resource_state.get('faElements', 0) if fa_count > 30: logger.info(f"{fa_count}個のFont Awesome要素があるため、追加待機...") time.sleep(min(1.0, fa_count / 100)) # 要素数に応じて待機（最大1秒） # コンテンツレンダリングのためのスクロール処理 - パフォーマンス改善 logger.info("Performing content rendering scroll...") total_height = driver.execute_script("return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight);") viewport_height = driver.execute_script("return window.innerHeight;") scrolls_needed = max(1, min(5, total_height // viewport_height)) # 最大5回までに制限 # スクロール処理の高速化 for i in range(scrolls_needed): scroll_pos = i * (viewport_height - 100) # 少しだけオーバーラップ driver.execute_script(f"window.scrollTo(0, {scroll_pos});") time.sleep(0.1) # 高速化のため待機時間短縮 # トップに戻る driver.execute_script("window.scrollTo(0, 0);") time.sleep(0.2) # 短い待機に変更 logger.info("Scroll rendering completed") # スクロールバーを非表示に driver.execute_script(""" document.documentElement.style.overflow = 'hidden'; document.body.style.overflow = 'hidden'; """) # ページの寸法を取得 dimensions = driver.execute_script(""" return { width: Math.max( document.documentElement.scrollWidth, document.documentElement.offsetWidth, document.documentElement.clientWidth, document.body ? document.body.scrollWidth : 0, document.body ? document.body.offsetWidth : 0, document.body ? document.body.clientWidth : 0 ), height: Math.max( document.documentElement.scrollHeight, document.documentElement.offsetHeight, document.documentElement.clientHeight, document.body ? document.body.scrollHeight : 0, document.body ? document.body.offsetHeight : 0, document.body ? document.body.clientHeight : 0 ) }; """) scroll_width = dimensions['width'] scroll_height = dimensions['height'] logger.info(f"Detected dimensions: width={scroll_width}, height={scroll_height}") # 最小/最大値の設定 scroll_width = max(scroll_width, 100) scroll_height = max(scroll_height, 100) scroll_width = min(scroll_width, 2000) scroll_height = min(scroll_height, 4000) # レイアウト安定化のための待機 time.sleep(2.0) # 高さに余白を追加 adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0)) adjusted_height = max(adjusted_height, scroll_height, 100) # ウィンドウサイズを調整 adjusted_width = scroll_width logger.info(f"Resizing window to: width={adjusted_width}, height={adjusted_height}") driver.set_window_size(adjusted_width, adjusted_height) time.sleep(0.5) # 短縮した待機時間 # スクリーンショット取得 logger.info("Taking screenshot...") png = driver.get_screenshot_as_png() logger.info("Screenshot taken successfully.") # PIL画像に変換 img = Image.open(BytesIO(png)) logger.info(f"Screenshot dimensions: {img.width}x{img.height}") # 余白トリミング - 最適化版を使用 if trim_whitespace: img = trim_image_whitespace(img, threshold=248, padding=20) logger.info(f"Trimmed dimensions: {img.width}x{img.height}") return img except Exception as e: logger.error(f"Error during screenshot generation: {e}", exc_info=True) # エラー時は小さい黒画像を返す return Image.new('RGB', (1, 1), color=(0, 0, 0)) finally: logger.info("Cleaning up...") # WebDriverプールに戻す if driver_from_pool: driver_pool.release_driver(driver) logger.info("Returned driver to pool") # 一時ファイル削除 if tmp_path and os.path.exists(tmp_path): try: os.remove(tmp_path) logger.info(f"Temporary file {tmp_path} removed.") except Exception as e: logger.error(f"Error removing temporary file {tmp_path}: {e}") # --- 並列処理を活用した新しい関数 --- def text_to_screenshot_parallel(text: str, extension_percentage: float, temperature: float = 0.5, trim_whitespace: bool = True, style: str = "standard") -> Image.Image: """テキストをGemini APIでHTMLに変換し、並列処理でスクリーンショットを生成する関数""" start_time = time.time() logger.info("並列処理によるテキスト→スクリーンショット生成を開始") try: # WebDriverと HTML生成を並列で実行 with ThreadPoolExecutor(max_workers=2) as executor: # Gemini APIリクエストタスク html_future = executor.submit( generate_html_from_text, text=text, temperature=temperature, style=style ) # WebDriver初期化タスク - プール使用 driver_future = executor.submit(driver_pool.get_driver) # 結果を取得 html_code = html_future.result() driver = driver_future.result() # ドライバーはプールから取得しているためフラグ設定 driver_from_pool = True # HTMLコードとドライバーが準備できたらスクリーンショット生成 logger.info(f"HTML生成完了：{len(html_code)}文字。スクリーンショット生成開始。") # レンダリング前にドライバーの初期設定 tmp_path = None try: # 一時ファイルにHTMLを保存 with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode='w', encoding='utf-8') as tmp_file: tmp_path = tmp_file.name tmp_file.write(html_code) logger.info(f"HTMLを一時ファイルに保存: {tmp_path}") # ウィンドウサイズ初期設定 initial_width = 1200 initial_height = 1000 driver.set_window_size(initial_width, initial_height) file_url = "file://" + tmp_path logger.info(f"ページに移動: {file_url}") driver.get(file_url) # ここからスクリーンショット生成ロジック（前の実装と同様） # ページ読み込み待機 - 動的な待機時間を実装 logger.info("body要素を待機...") WebDriverWait(driver, 10).until( EC.presence_of_element_located((By.TAG_NAME, "body")) ) logger.info("body要素を検出。リソース読み込みを待機...") # リソース読み込みの動的待機 - 最適化 max_wait = 3 # 最大待機時間（秒） wait_increment = 0.2 # 確認間隔 wait_time = 0 while wait_time < max_wait: resource_state = driver.execute_script(""" return { complete: document.readyState === 'complete', imgCount: document.images.length, imgLoaded: Array.from(document.images).filter(img => img.complete).length, faElements: document.querySelectorAll('.fa, .fas, .far, .fab, [class*="fa-"]').length }; """) # ドキュメント完了かつ画像が読み込まれている場合、待機終了 if resource_state['complete'] and (resource_state['imgCount'] == 0 or resource_state['imgLoaded'] == resource_state['imgCount']): logger.info(f"リソース読み込み完了: {resource_state}") break time.sleep(wait_increment) wait_time += wait_increment # Font Awesome要素が多い場合は追加待機 fa_count = resource_state.get('faElements', 0) if fa_count > 30: logger.info(f"{fa_count}個のFont Awesome要素があるため、追加待機...") time.sleep(min(1.0, fa_count / 100)) # 要素数に応じて待機（最大0.5秒） # コンテンツレンダリングのための簡易スクロール driver.execute_script("window.scrollTo(0, document.body.scrollHeight);") time.sleep(0.2) driver.execute_script("window.scrollTo(0, 0);") time.sleep(0.2) # スクロールバーを非表示に driver.execute_script(""" document.documentElement.style.overflow = 'hidden'; document.body.style.overflow = 'hidden'; """) # ページの寸法を取得 dimensions = driver.execute_script(""" return { width: Math.max( document.documentElement.scrollWidth, document.documentElement.offsetWidth, document.documentElement.clientWidth, document.body ? document.body.scrollWidth : 0, document.body ? document.body.offsetWidth : 0, document.body ? document.body.clientWidth : 0 ), height: Math.max( document.documentElement.scrollHeight, document.documentElement.offsetHeight, document.documentElement.clientHeight, document.body ? document.body.scrollHeight : 0, document.body ? document.body.offsetHeight : 0, document.body ? document.body.clientHeight : 0 ) }; """) scroll_width = dimensions['width'] scroll_height = dimensions['height'] # 最小/最大値の設定 scroll_width = max(scroll_width, 100) scroll_height = max(scroll_height, 100) scroll_width = min(scroll_width, 2000) scroll_height = min(scroll_height, 4000) # 高さに余白を追加 adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0)) adjusted_height = max(adjusted_height, scroll_height, 100) # ウィンドウサイズを調整 driver.set_window_size(scroll_width, adjusted_height) time.sleep(0.2) # スクリーンショット取得 logger.info("スクリーンショットを撮影...") png = driver.get_screenshot_as_png() # PIL画像に変換 img = Image.open(BytesIO(png)) logger.info(f"スクリーンショットサイズ: {img.width}x{img.height}") # 余白トリミング if trim_whitespace: img = trim_image_whitespace(img, threshold=248, padding=20) logger.info(f"トリミング後のサイズ: {img.width}x{img.height}") elapsed = time.time() - start_time logger.info(f"並列処理による生成完了。所要時間: {elapsed:.2f}秒") return img except Exception as e: logger.error(f"スクリーンショット生成中にエラー: {e}", exc_info=True) return Image.new('RGB', (1, 1), color=(0, 0, 0)) finally: # WebDriverプールに戻す if driver_from_pool: driver_pool.release_driver(driver) # 一時ファイル削除 if tmp_path and os.path.exists(tmp_path): try: os.remove(tmp_path) except Exception as e: logger.error(f"一時ファイル削除エラー: {e}") except Exception as e: logger.error(f"並列処理中のエラー: {e}", exc_info=True) return Image.new('RGB', (1, 1), color=(0, 0, 0)) # エラー時は黒画像 # 従来の非並列版も残す（互換性のため） def text_to_screenshot(text: str, extension_percentage: float, temperature: float = 0.3, trim_whitespace: bool = True, style: str = "standard") -> Image.Image: """テキストをGemini APIでHTMLに変換し、スクリーンショットを生成する統合関数（レガシー版）""" # 並列処理版を呼び出す return text_to_screenshot_parallel(text, extension_percentage, temperature, trim_whitespace, style) # --- FastAPI Setup --- app = FastAPI() # CORS設定を追加 app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 静的ファイルのサービング設定 # Gradioのディレクトリを探索してアセットを見つける gradio_dir = os.path.dirname(gr.__file__) logger.info(f"Gradio version: {gr.__version__}") logger.info(f"Gradio directory: {gradio_dir}") # 基本的な静的ファイルディレクトリをマウント static_dir = os.path.join(gradio_dir, "templates", "frontend", "static") if os.path.exists(static_dir): logger.info(f"Mounting static directory: {static_dir}") app.mount("/static", StaticFiles(directory=static_dir), name="static") # _appディレクトリを探す（新しいSvelteKitベースのフロントエンド用） app_dir = os.path.join(gradio_dir, "templates", "frontend", "_app") if os.path.exists(app_dir): logger.info(f"Mounting _app directory: {app_dir}") app.mount("/_app", StaticFiles(directory=app_dir), name="_app") # assetsディレクトリを探す assets_dir = os.path.join(gradio_dir, "templates", "frontend", "assets") if os.path.exists(assets_dir): logger.info(f"Mounting assets directory: {assets_dir}") app.mount("/assets", StaticFiles(directory=assets_dir), name="assets") # cdnディレクトリがあれば追加 cdn_dir = os.path.join(gradio_dir, "templates", "cdn") if os.path.exists(cdn_dir): logger.info(f"Mounting cdn directory: {cdn_dir}") app.mount("/cdn", StaticFiles(directory=cdn_dir), name="cdn") # API Endpoint for screenshot generation @app.post("/api/screenshot", response_class=StreamingResponse, tags=["Screenshot"], summary="Render HTML to Full Page Screenshot", description="Takes HTML code and an optional vertical extension percentage, renders it using a headless browser, and returns the full-page screenshot as a PNG image.") async def api_render_screenshot(request: ScreenshotRequest): """ API endpoint to render HTML and return a screenshot. """ try: logger.info(f"API request received. Extension: {request.extension_percentage}%") # Run the blocking Selenium code (now using the pooled version) pil_image = render_fullpage_screenshot( request.html_code, request.extension_percentage, request.trim_whitespace ) if pil_image.size == (1, 1): logger.error("Screenshot generation failed, returning 1x1 error image.") # Optionally return a proper error response instead of 1x1 image # raise HTTPException(status_code=500, detail="Failed to generate screenshot") # Convert PIL Image to PNG bytes img_byte_arr = BytesIO() pil_image.save(img_byte_arr, format='PNG') img_byte_arr.seek(0) # Go to the start of the BytesIO buffer logger.info("Returning screenshot as PNG stream.") return StreamingResponse(img_byte_arr, media_type="image/png") except Exception as e: logger.error(f"API Error: {e}", exc_info=True) raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}") # --- 新しいGemini API連携エンドポイント（並列処理版） --- @app.post("/api/text-to-screenshot", response_class=StreamingResponse, tags=["Screenshot", "Gemini"], summary="テキストからインフォグラフィックを生成", description="テキストをGemini APIを使ってHTMLインフォグラフィックに変換し、スクリーンショットとして返します。") async def api_text_to_screenshot(request: GeminiRequest): """ テキストからHTMLインフォグラフィックを生成してスクリーンショットを返すAPIエンドポイント """ try: logger.info(f"テキスト→スクリーンショットAPIリクエスト受信。テキスト長さ: {len(request.text)}, " f"拡張率: {request.extension_percentage}%, 温度: {request.temperature}, " f"スタイル: {request.style}") # 並列処理版を使用 pil_image = text_to_screenshot_parallel( request.text, request.extension_percentage, request.temperature, request.trim_whitespace, request.style ) if pil_image.size == (1, 1): logger.error("スクリーンショット生成に失敗しました。1x1エラー画像を返します。") # PIL画像をPNGバイトに変換 img_byte_arr = BytesIO() pil_image.save(img_byte_arr, format='PNG') img_byte_arr.seek(0) # BytesIOバッファの先頭に戻る logger.info("スクリーンショットをPNGストリームとして返します。") return StreamingResponse(img_byte_arr, media_type="image/png") except Exception as e: logger.error(f"API Error: {e}", exc_info=True) raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}") # --- Gradio Interface Definition --- # 入力モードの選択用Radioコンポーネント def process_input(input_mode, input_text, extension_percentage, temperature, trim_whitespace, style): """入力モードに応じて適切な処理を行う""" if input_mode == "HTML入力": # HTMLモードの場合は既存の処理（スタイルは使わない） return render_fullpage_screenshot(input_text, extension_percentage, trim_whitespace) else: # テキスト入力モードの場合はGemini APIを使用（並列処理版） return text_to_screenshot_parallel(input_text, extension_percentage, temperature, trim_whitespace, style) # Gradio UIの定義 with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr.themes.Base()) as iface: gr.Markdown("# HTMLビューア & テキスト→インフォグラフィック変換") gr.Markdown("HTMLコードをレンダリングするか、テキストをGemini APIでインフォグラフィックに変換して画像として取得します。") gr.Markdown("**パフォーマンス向上版**: 並列処理と最適化により処理時間を短縮しています") with gr.Row(): input_mode = gr.Radio( ["HTML入力", "テキスト入力"], label="入力モード", value="HTML入力" ) # 共用のテキストボックス input_text = gr.Textbox( lines=15, label="入力", placeholder="HTMLコードまたはテキストを入力してください。入力モードに応じて処理されます。" ) with gr.Row(): with gr.Column(scale=1): # スタイル選択ドロップダウン style_dropdown = gr.Dropdown( choices=["standard", "cute", "resort", "cool", "dental", "school", "KOKUGO"], value="standard", label="デザインスタイル", info="テキスト→HTML変換時のデザインテーマを選択します", visible=False # テキスト入力モードの時だけ表示 ) with gr.Column(scale=2): extension_percentage = gr.Slider( minimum=0, maximum=30, step=1.0, value=10, # デフォルト値10% label="上下高さ拡張率（%）" ) # 温度調整スライダー（テキストモード時のみ表示） temperature = gr.Slider( minimum=0.0, maximum=1.0, step=0.1, value=0.5, # デフォルト値を0.5に設定 label="生成時の温度（低い=一貫性高、高い=創造性高）", visible=False # 最初は非表示 ) # 余白トリミングオプション trim_whitespace = gr.Checkbox( label="余白を自動トリミング", value=True, info="生成される画像から余分な空白領域を自動的に削除します" ) submit_btn = gr.Button("生成") output_image = gr.Image(type="pil", label="ページ全体のスクリーンショット") # 入力モード変更時のイベント処理（テキストモード時のみ温度スライダーとスタイルドロップダウンを表示） def update_controls_visibility(mode): # Gradio 4.x用のアップデート方法 is_text_mode = mode == "テキスト入力" return [ gr.update(visible=is_text_mode), # temperature gr.update(visible=is_text_mode), # style_dropdown ] input_mode.change( fn=update_controls_visibility, inputs=input_mode, outputs=[temperature, style_dropdown] ) # 生成ボタンクリック時のイベント処理 submit_btn.click( fn=process_input, inputs=[input_mode, input_text, extension_percentage, temperature, trim_whitespace, style_dropdown], outputs=output_image ) # 環境変数情報を表示 gemini_model = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro") gr.Markdown(f""" ## APIエンドポイント - `/api/screenshot` - HTMLコードからスクリーンショットを生成 - `/api/text-to-screenshot` - テキストからインフォグラフィックスクリーンショットを生成 ## 設定情報 - 使用モデル: {gemini_model} (環境変数 GEMINI_MODEL で変更可能) - 対応スタイル: standard, cute, resort, cool, dental, school, KOKUGO - WebDriverプール最大数: {driver_pool.max_drivers} (環境変数 MAX_WEBDRIVERS で変更可能) """) # --- Mount Gradio App onto FastAPI --- app = gr.mount_gradio_app(app, iface, path="/") # --- Run with Uvicorn (for local testing) --- if __name__ == "__main__": import uvicorn logger.info("Starting Uvicorn server for local development...") uvicorn.run(app, host="0.0.0.0", port=7860) # アプリケーション終了時にWebDriverプールをクリーンアップ import atexit atexit.register(driver_pool.close_all)