Spaces:
Paused
Paused
| 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 | |
| from huggingface_hub import hf_hub_download # 追加: HuggingFace Hubからファイルを直接ダウンロード | |
| # 正しいGemini関連のインポート | |
| import google.generativeai as genai | |
| # ロギング設定 | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # --- 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レイアウトを改善するCSSを追加""" | |
| # CSSを追加 | |
| fa_fix_css = """ | |
| <style> | |
| /* Font Awesomeアイコンのレイアウト修正 */ | |
| [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> | |
| """ | |
| # headタグがある場合はその中に追加 | |
| if '<head>' in html_code: | |
| return html_code.replace('</head>', f'{fa_fix_css}</head>') | |
| # HTMLタグがある場合はその後に追加 | |
| elif '<html' in html_code: | |
| head_end = html_code.find('</head>') | |
| if head_end > 0: | |
| return html_code[:head_end] + fa_fix_css + html_code[head_end:] | |
| else: | |
| body_start = html_code.find('<body') | |
| if body_start > 0: | |
| return html_code[:body_start] + f'<head>{fa_fix_css}</head>' + html_code[body_start:] | |
| # どちらもない場合は先頭に追加 | |
| return f'<html><head>{fa_fix_css}</head>' + html_code + '</html>' | |
| def load_system_instruction(style="standard"): | |
| """ | |
| 指定されたスタイルのシステムインストラクションを読み込む | |
| Args: | |
| style: 使用するスタイル名 (standard, cute, resort, cool, dental) | |
| Returns: | |
| 読み込まれたシステムインストラクション | |
| """ | |
| try: | |
| # 有効なスタイル一覧 | |
| valid_styles = ["standard", "cute", "resort", "cool", "dental"] | |
| # スタイルの検証 | |
| 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.3, 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}") | |
| # 画像から余分な空白領域をトリミングする関数 | |
| def trim_image_whitespace(image, threshold=250, padding=10): | |
| """ | |
| 画像から余分な白い空白をトリミングする | |
| Args: | |
| image: PIL.Image - 入力画像 | |
| threshold: int - どの明るさ以上を空白と判断するか (0-255) | |
| padding: int - トリミング後に残す余白のピクセル数 | |
| Returns: | |
| トリミングされたPIL.Image | |
| """ | |
| # グレースケールに変換 | |
| gray = image.convert('L') | |
| # ピクセルデータを配列として取得 | |
| data = gray.getdata() | |
| width, height = gray.size | |
| # 有効範囲を見つける | |
| min_x, min_y = width, height | |
| max_x = max_y = 0 | |
| # ピクセルデータを2次元配列に変換して処理 | |
| pixels = list(data) | |
| pixels = [pixels[i * width:(i + 1) * width] for i in range(height)] | |
| # 各行をスキャンして非空白ピクセルを見つける | |
| for y in range(height): | |
| for x in range(width): | |
| if pixels[y][x] < threshold: # 非空白ピクセル | |
| min_x = min(min_x, x) | |
| min_y = min(min_y, y) | |
| max_x = max(max_x, x) | |
| max_y = max(max_y, y) | |
| # 境界外のトリミングの場合はエラー | |
| if min_x > max_x or min_y > max_y: | |
| logger.warning("トリミング領域が見つかりません。元の画像を返します。") | |
| return image | |
| # パディングを追加 | |
| min_x = max(0, min_x - padding) | |
| min_y = max(0, min_y - padding) | |
| max_x = min(width - 1, max_x + padding) | |
| max_y = min(height - 1, max_y + padding) | |
| # 画像をトリミング | |
| trimmed = image.crop((min_x, min_y, max_x + 1, max_y + 1)) | |
| logger.info(f"画像をトリミングしました: 元サイズ {width}x{height} → トリミング後 {trimmed.width}x{trimmed.height}") | |
| return trimmed | |
| # 非同期スクリプトを使わず、同期的なスクリプトのみ使用する改善版 | |
| def render_fullpage_screenshot(html_code: str, extension_percentage: float = 6.0, | |
| trim_whitespace: bool = True) -> Image.Image: | |
| """ | |
| Renders HTML code to a full-page screenshot using Selenium. | |
| 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. | |
| Returns: | |
| A PIL Image object of the screenshot. | |
| """ | |
| tmp_path = None | |
| driver = None | |
| # 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}") | |
| return Image.new('RGB', (1, 1), color=(0, 0, 0)) | |
| # 2) Headless Chrome(Chromium) options | |
| 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"Using CHROMEDRIVER_PATH: {webdriver_path}") | |
| service = webdriver.ChromeService(executable_path=webdriver_path) | |
| else: | |
| logger.info("CHROMEDRIVER_PATH not set or invalid, using default PATH lookup.") | |
| service = None # Use default behavior | |
| try: | |
| logger.info("Initializing WebDriver...") | |
| if service: | |
| driver = webdriver.Chrome(service=service, options=options) | |
| else: | |
| driver = webdriver.Chrome(options=options) | |
| logger.info("WebDriver initialized.") | |
| # 3) 初期ウィンドウサイズを設定 | |
| 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) | |
| # 4) ページ読み込み待機 | |
| logger.info("Waiting for body element...") | |
| WebDriverWait(driver, 15).until( | |
| EC.presence_of_element_located((By.TAG_NAME, "body")) | |
| ) | |
| logger.info("Body element found. Waiting for resource loading...") | |
| # 5) 基本的なリソース読み込み待機 - タイムアウト回避 | |
| time.sleep(3) | |
| # Font Awesome読み込み確認 - 非同期を使わない | |
| logger.info("Checking for Font Awesome resources...") | |
| fa_count = driver.execute_script(""" | |
| var icons = document.querySelectorAll('.fa, .fas, .far, .fab, [class*="fa-"]'); | |
| return icons.length; | |
| """) | |
| logger.info(f"Found {fa_count} Font Awesome elements") | |
| # リソース読み込み状態を確認 | |
| doc_ready = driver.execute_script("return document.readyState;") | |
| logger.info(f"Document ready state: {doc_ready}") | |
| # Font Awesomeが多い場合は追加待機 | |
| if fa_count > 50: | |
| logger.info("Many Font Awesome icons detected, waiting additional time") | |
| time.sleep(2) | |
| # 6) コンテンツレンダリングのためのスクロール処理 - 同期的に実行 | |
| 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, total_height // viewport_height) | |
| for i in range(scrolls_needed + 1): | |
| scroll_pos = i * (viewport_height - 200) # オーバーラップさせる | |
| driver.execute_script(f"window.scrollTo(0, {scroll_pos});") | |
| time.sleep(0.2) # 短い待機 | |
| # トップに戻る | |
| driver.execute_script("window.scrollTo(0, 0);") | |
| time.sleep(0.5) | |
| logger.info("Scroll rendering completed") | |
| # 7) スクロールバーを非表示に | |
| driver.execute_script(""" | |
| document.documentElement.style.overflow = 'hidden'; | |
| document.body.style.overflow = 'hidden'; | |
| """) | |
| logger.info("Scrollbars hidden") | |
| # 8) ページの寸法を取得 | |
| 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}") | |
| # 再検証 - 短いスクロールで再確認 | |
| driver.execute_script("window.scrollTo(0, document.body.scrollHeight);") | |
| time.sleep(0.5) | |
| driver.execute_script("window.scrollTo(0, 0);") | |
| time.sleep(0.5) | |
| dimensions_after = driver.execute_script("return {height: Math.max(document.documentElement.scrollHeight, document.body.scrollHeight)};") | |
| scroll_height = max(scroll_height, dimensions_after['height']) | |
| logger.info(f"After scroll check, 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) | |
| # 9) レイアウト安定化のための単純な待機 - タイムアウト回避 | |
| logger.info("Waiting for layout stabilization...") | |
| time.sleep(2) | |
| # 10) 高さに余白を追加 | |
| adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0)) | |
| adjusted_height = max(adjusted_height, scroll_height, 100) | |
| logger.info(f"Adjusted height calculated: {adjusted_height} (extension: {extension_percentage}%)") | |
| # 11) ウィンドウサイズを調整 | |
| 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(1) | |
| # リソース状態を確認 - 同期的スクリプト | |
| resource_state = driver.execute_script(""" | |
| return { | |
| readyState: document.readyState, | |
| resourcesComplete: !document.querySelector('img:not([complete])') && | |
| !document.querySelector('link[rel="stylesheet"]:not([loaded])') | |
| }; | |
| """) | |
| logger.info(f"Resource state: {resource_state}") | |
| if resource_state['readyState'] != 'complete': | |
| logger.info("Document still loading, waiting additional time...") | |
| time.sleep(1) | |
| # トップにスクロール | |
| driver.execute_script("window.scrollTo(0, 0);") | |
| time.sleep(0.5) | |
| logger.info("Scrolled to top.") | |
| # 12) スクリーンショット取得 | |
| 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 a small black image on error | |
| return Image.new('RGB', (1, 1), color=(0, 0, 0)) | |
| finally: | |
| logger.info("Cleaning up...") | |
| if driver: | |
| try: | |
| driver.quit() | |
| logger.info("WebDriver quit successfully.") | |
| except Exception as e: | |
| logger.error(f"Error quitting WebDriver: {e}", exc_info=True) | |
| 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}", exc_info=True) | |
| # --- Geminiを使った新しい関数 --- | |
| 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に変換し、スクリーンショットを生成する統合関数""" | |
| try: | |
| # 1. テキストからHTMLを生成（温度パラメータとスタイルも渡す） | |
| html_code = generate_html_from_text(text, temperature, style) | |
| # 2. HTMLからスクリーンショットを生成 | |
| return render_fullpage_screenshot(html_code, extension_percentage, trim_whitespace) | |
| except Exception as e: | |
| logger.error(f"テキストからスクリーンショット生成中にエラーが発生: {e}", exc_info=True) | |
| return Image.new('RGB', (1, 1), color=(0, 0, 0)) # エラー時は黒画像 | |
| # --- 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 | |
| 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 in a separate thread (FastAPI handles this) | |
| 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連携エンドポイント --- | |
| 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}") | |
| # テキストからHTMLを生成してスクリーンショットを作成（温度パラメータとスタイルも渡す） | |
| pil_image = text_to_screenshot( | |
| request.text, | |
| request.extension_percentage, | |
| request.temperature, | |
| request.trim_whitespace, | |
| request.style | |
| ) | |
| if pil_image.size == (1, 1): | |
| logger.error("スクリーンショット生成に失敗しました。1x1エラー画像を返します。") | |
| # raise HTTPException(status_code=500, detail="スクリーンショット生成に失敗しました") | |
| # 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(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でインフォグラフィックに変換して画像として取得します。") | |
| 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"], | |
| 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 | |
| """) | |
| # --- 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) |