File size: 13,790 Bytes
d9bd4cb
03a6dcc
d9bd4cb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
03a6dcc
d9bd4cb
03a6dcc
 
 
 
d9bd4cb
 
 
 
 
03a6dcc
d9bd4cb
 
03a6dcc
 
 
 
d9bd4cb
 
 
03a6dcc
 
 
d9bd4cb
03a6dcc
d9bd4cb
 
 
03a6dcc
 
 
 
 
 
 
 
 
d9bd4cb
 
 
 
 
 
 
 
03a6dcc
 
 
aeb0d3b
d9bd4cb
 
03a6dcc
 
 
 
 
 
 
 
d9bd4cb
03a6dcc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d9bd4cb
03a6dcc
d333263
03a6dcc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
df45e87
aeb0d3b
03a6dcc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7215bb7
03a6dcc
 
 
 
 
 
 
 
 
 
aeb0d3b
 
 
 
03a6dcc
 
 
aeb0d3b
03a6dcc
 
 
aeb0d3b
03a6dcc
 
 
 
 
 
 
 
 
 
 
aeb0d3b
 
03a6dcc
 
 
 
 
 
 
 
 
aeb0d3b
 
 
03a6dcc
aeb0d3b
03a6dcc
 
aeb0d3b
 
03a6dcc
aeb0d3b
03a6dcc
aeb0d3b
03a6dcc
 
 
 
aeb0d3b
03a6dcc
aeb0d3b
 
 
03a6dcc
 
aeb0d3b
 
03a6dcc
aeb0d3b
03a6dcc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2917f52
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
03a6dcc
2917f52
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
aeb0d3b
d9bd4cb
2917f52
d9bd4cb
 
03a6dcc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
import gradio as gr
from fastapi import FastAPI, HTTPException
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

# --- Gemini SDK (v1.x) ---------------------------------
from google import genai
from google.genai import types
# -------------------------------------------------------

# ロギング設定
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ---------- データモデル ----------
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"

# ---------- ユーティリティ ----------
def enhance_font_awesome_layout(html_code):
    fa_fix_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+.fab{
        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_fix_css}</head>')
    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:]
        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"):
    valid_styles = ["standard", "cute", "resort", "cool", "dental"]
    if style not in valid_styles:
        logger.warning(f"無効なスタイル '{style}'。'standard' を使用")
        style = "standard"
    local_path = os.path.join(os.path.dirname(__file__), style, "prompt.txt")
    if os.path.exists(local_path):
        with open(local_path, encoding="utf-8") as f:
            return f.read()
    try:
        file_path = hf_hub_download(
            repo_id="tomo2chin2/GURAREKOstlyle",
            filename=f"{style}/prompt.txt",
            repo_type="dataset"
        )
        with open(file_path, encoding="utf-8") as f:
            return f.read()
    except Exception:
        file_path = hf_hub_download(
            repo_id="tomo2chin2/GURAREKOstlyle",
            filename="prompt.txt",
            repo_type="dataset"
        )
        with open(file_path, encoding="utf-8") as f:
            return f.read()

# ---------- Gemini HTML 生成 ----------
def generate_html_from_text(text, temperature=0.3, style="standard"):
    """Gemini で HTML を生成。2.5 Flash Preview の場合は thinking_off"""
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY が未設定")
    model_name = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro")

    client = genai.Client(api_key=api_key)
    system_instruction = load_system_instruction(style)
    prompt = f"{system_instruction}\n\n{text}"

    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"},
    ]

    # ---- モデル分岐 ----
    if model_name == "gemini-2.5-flash-preview-04-17":
        generation_cfg = types.GenerateContentConfig(
            temperature=temperature,
            top_p=0.7,
            top_k=20,
            max_output_tokens=8192,
            candidate_count=1,
            thinking_config=types.ThinkingConfig(thinking_budget=0)  # thinking OFF
        )
    else:
        generation_cfg = types.GenerateContentConfig(
            temperature=temperature,
            top_p=0.7,
            top_k=20,
            max_output_tokens=8192,
            candidate_count=1,
        )

    response = client.models.generate_content(
        model=model_name,
        contents=prompt,
        config=generation_cfg,
        safety_settings=safety_settings
    )
    raw = response.text

    # ```html ... ``` 抽出
    start = raw.find("```html")
    end = raw.rfind("```")
    if start != -1 and end != -1 and start < end:
        html_code = raw[start + 7:end].strip()
        return enhance_font_awesome_layout(html_code)
    logger.warning("```html``` ブロックが見つからず全文返却")
    return raw

# ---------- 画像トリミング ----------
def trim_image_whitespace(image, threshold=250, padding=10):
    gray = image.convert("L")
    data = list(gray.getdata())
    w, h = gray.size
    pixels = [data[i * w:(i + 1) * w] for i in range(h)]
    min_x, min_y, max_x, max_y = w, h, 0, 0
    for y in range(h):
        for x in range(w):
            if pixels[y][x] < threshold:
                min_x, min_y = min(min_x, x), min(min_y, y)
                max_x, max_y = max(max_x, x), max(max_y, y)
    if min_x > max_x or min_y > max_y:
        return image
    min_x = max(0, min_x - padding)
    min_y = max(0, min_y - padding)
    max_x = min(w - 1, max_x + padding)
    max_y = min(h - 1, max_y + padding)
    return image.crop((min_x, min_y, max_x + 1, max_y + 1))

# ---------- HTML → スクリーンショット(Selenium) ----------
def render_fullpage_screenshot(html_code, extension_percentage=6.0, trim_whitespace=True):
    tmp_path, driver = None, None
    try:
        with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode="w", encoding="utf-8") as f:
            tmp_path = f.name
            f.write(html_code)
        options = Options()
        options.add_argument("--headless")
        options.add_argument("--no-sandbox")
        options.add_argument("--disable-dev-shm-usage")
        driver = webdriver.Chrome(options=options)
        driver.set_window_size(1200, 1000)
        driver.get("file://" + tmp_path)
        WebDriverWait(driver, 15).until(EC.presence_of_element_located((By.TAG_NAME, "body")))
        time.sleep(3)

        total_height = driver.execute_script("return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)")
        viewport = driver.execute_script("return window.innerHeight")
        for i in range(max(1, total_height // viewport) + 1):
            driver.execute_script(f"window.scrollTo(0, {i * (viewport - 200)})")
            time.sleep(0.2)
        driver.execute_script("window.scrollTo(0, 0)")
        driver.execute_script("document.documentElement.style.overflow='hidden';document.body.style.overflow='hidden'")
        dims = driver.execute_script("return {w:document.documentElement.scrollWidth,h:document.documentElement.scrollHeight}")
        driver.set_window_size(dims["w"], int(dims["h"] * (1 + extension_percentage / 100)))
        time.sleep(1)
        png = driver.get_screenshot_as_png()
        img = Image.open(BytesIO(png))
        return trim_image_whitespace(img, 248, 20) if trim_whitespace else img
    except Exception as e:
        logger.error(f"Selenium error: {e}", exc_info=True)
        return Image.new("RGB", (1, 1))
    finally:
        if driver:
            driver.quit()
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)

# ---------- 統合 ----------
def text_to_screenshot(text, ext, temp=0.3, trim=True, style="standard"):
    html = generate_html_from_text(text, temp, style)
    return render_fullpage_screenshot(html, ext, trim)

# ---------- FastAPI ----------
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], allow_credentials=True,
    allow_methods=["*"], allow_headers=["*"],
)

# Gradio 静的ファイル
gradio_dir = os.path.dirname(gr.__file__)
app.mount("/static", StaticFiles(directory=os.path.join(gradio_dir, "templates/frontend/static")), name="static")

# ---------- API ----------
@app.post("/api/screenshot", response_class=StreamingResponse)
async def api_render(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)
async def api_text_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")

# ---------- Gradio ----------
def process_input(mode, txt, ext, temp, trim, style):
    if mode == "HTML入力":
        return render_fullpage_screenshot(txt, ext, trim)
    return text_to_screenshot(txt, ext, temp, trim, style)

# ---------- Gradio UI ----------
with gr.Blocks(title="Full Page Screenshot (テキスト変換対応)", theme=gr.themes.Base()) as iface:
    # 見出し
    gr.Markdown(
        "<h1 style='text-align:center;margin:0.2em 0'>HTMLビューア & テキスト→インフォグラフィック変換</h1>",
        elem_id="title",
        inline=True,
    )
    gr.Markdown(
        "HTML を直接レンダリングするか、テキストを Gemini API でインフォグラフィックに変換して画像取得できます。",
        elem_id="subtitle",
    )

    # === 入力モード選択 ===
    with gr.Row():
        input_mode = gr.Radio(
            ["HTML入力", "テキスト入力"],
            value="HTML入力",
            label="入力モード",
        )

    # === 共通入力テキスト ===
    input_text = gr.Textbox(
        lines=15,
        label="入力",
        placeholder="HTMLコードまたはテキストを入力してください(モードに応じて処理)。",
    )

    # === スタイル + スライダー類を 2 カラムで ===
    with gr.Row():
        # 左:スタイル
        with gr.Column(scale=1, min_width=180):
            style_dropdown = gr.Dropdown(
                choices=["standard", "cute", "resort", "cool", "dental"],
                value="standard",
                label="デザインスタイル",
                info="テキスト→HTML 変換時のテーマ",
                visible=False,              # 初期は非表示
            )
        # 右:拡張率 & 温度
        with gr.Column(scale=3):
            extension_percentage = gr.Slider(
                0, 30, value=10, step=1,
                label="上下高さ拡張率(%)"
            )
            temperature = gr.Slider(
                0.0, 1.0, value=0.5, step=0.1,
                label="生成時の温度(低い=一貫性高、高い=創造性高)",
                visible=False,              # 初期は非表示
            )

    # === トリミング & ボタン ===
    trim_whitespace = gr.Checkbox(
        value=True,
        label="余白を自動トリミング",
    )
    submit_btn = gr.Button("生成", variant="primary", size="lg")

    # === 出力 ===
    output_image = gr.Image(
        type="pil",
        label="ページ全体のスクリーンショット",
        show_label=True,
        show_download_button=True,
    )

    # ---------- 変化イベント ----------
    def toggle_controls(mode):
        """テキストモードのときだけ温度とスタイルを表示"""
        is_text = mode == "テキスト入力"
        return (
            gr.update(visible=is_text),   # temperature
            gr.update(visible=is_text),   # style_dropdown
        )

    input_mode.change(
        fn=toggle_controls,
        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,
    )

    # 下部にモデル表示
    gr.Markdown(
        f"""
**使用モデル** : `{os.getenv("GEMINI_MODEL", "gemini-1.5-pro")}`  
**API** : `/api/screenshot`  /  `/api/text-to-screenshot`
""",
        elem_id="footnote",
    )

# --- Gradio を FastAPI にマウント ---
app = gr.mount_gradio_app(app, iface, path="/")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)