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"", text, re.IGNORECASE) end_match = re.search(r"", 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}")