| import os |
| import re |
| from io import BytesIO |
| from PIL import Image |
| from playwright.sync_api import sync_playwright |
|
|
|
|
| def extract_clean_html(text): |
| """ |
| 清理模型输出，尽量提取完整 HTML。 |
| """ |
| text = text.replace("```html", "").replace("```", "") |
|
|
| start_match = re.search(r"<!DOCTYPE html>", text, re.IGNORECASE) |
| end_match = re.search(r"</html>", text, re.IGNORECASE) |
|
|
| if start_match and end_match: |
| start_idx = start_match.start() |
| end_idx = end_match.end() |
| if end_idx > start_idx: |
| return text[start_idx:end_idx] |
|
|
| return text.strip() |
|
|
|
|
| def render_html_to_image(html, width=1080, height=2400): |
| """ |
| 用 Playwright 把 HTML 渲染成 PIL.Image。 |
| """ |
| with sync_playwright() as p: |
| browser = p.chromium.launch(headless=True) |
| page = browser.new_page(viewport={"width": width, "height": height}) |
| page.set_content(html, wait_until="domcontentloaded") |
| screenshot_bytes = page.screenshot(full_page=False) |
| browser.close() |
|
|
| return Image.open(BytesIO(screenshot_bytes)).convert("RGB") |
|
|
|
|
| def save_demo_outputs(output_dir, hinted_image, html, rendered_image=None): |
| """ |
| 保存 demo 输出： |
| - input_with_hint.png |
| - predicted_next_ui.html |
| - predicted_next_ui.png |
| """ |
| os.makedirs(output_dir, exist_ok=True) |
|
|
| hinted_path = os.path.join(output_dir, "input_with_hint.png") |
| html_path = os.path.join(output_dir, "predicted_next_ui.html") |
| rendered_path = os.path.join(output_dir, "predicted_next_ui.png") |
|
|
| hinted_image.save(hinted_path) |
|
|
| with open(html_path, "w", encoding="utf-8") as f: |
| f.write(html) |
|
|
| if rendered_image is not None: |
| rendered_image.save(rendered_path) |
|
|
| print(f"Saved hinted image to: {hinted_path}") |
| print(f"Saved HTML to: {html_path}") |
| if rendered_image is not None: |
| print(f"Saved rendered image to: {rendered_path}") |
