Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # must come before torch | |
| import html as html_lib | |
| import re | |
| import tempfile | |
| import time | |
| from threading import Thread | |
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| from transformers import AutoModelForImageTextToText, AutoProcessor, TextIteratorStreamer | |
| MODEL_ID = "microsoft/AesCode-8B" | |
| processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| model = AutoModelForImageTextToText.from_pretrained( | |
| MODEL_ID, dtype=torch.bfloat16, attn_implementation="sdpa" | |
| ).to("cuda") | |
| model.eval() | |
| # --------------------------------------------------------------------------- | |
| # Prompting — copied from github.com/microsoft/AesCode inference/prompts.py | |
| # --------------------------------------------------------------------------- | |
| SYSTEM_PROMPT = ( | |
| "You create polished standalone HTML/CSS infographics, slides, and editorial layouts. " | |
| "You receive a generation prompt and a reference image. " | |
| "The generation prompt is the primary contract: follow it exactly wherever it specifies required text, charts, tables, or data. " | |
| "The reference image is visual guidance only: reuse strong ideas in hierarchy, palette, density, composition, and typography, but do not treat it as a screenshot to copy. " | |
| "Its visible text, numbers, and labels may be incomplete, inaccurate, or stylistically decorative, so do not trust them over the prompt or inline data contracts. " | |
| "HTML output constraints: " | |
| "(1) Produce exactly one complete <!DOCTYPE html> document with all CSS in a single inline <style> block. " | |
| "(2) The root artboard element (e.g. .poster, .artboard, or body > div) MUST declare explicit width and height in pixels matching the canvas size specified in the prompt. Use a single top-level container with fixed px dimensions — never 100vh, 100%, or viewport-relative units for the root canvas. " | |
| "(3) External resources are limited to https://placehold.co for content images and https://cdn.jsdelivr.net/npm/echarts@5.5.0/dist/echarts.min.js for charts. " | |
| "(4) ECharts series.data must be literal arrays, never computed at runtime. " | |
| "(5) Inline '(Data contract [...])' clauses contain exact required values and must be copied verbatim into ECharts series.data or visible HTML table cells. " | |
| "(6) Tables must use real <table>/<tr>/<td> elements with visible text. " | |
| "(7) The root artboard must use overflow:hidden, and all content must fit inside the declared canvas size. " | |
| "(8) Each chart, table, title, heading, and footer must remain a separate DOM element, not merged into a decorative background or shared node. " | |
| "When refining a previous implementation, you also receive a low-resolution screenshot of that implementation; compare it against the reference image to identify visual discrepancies and fix the HTML accordingly." | |
| ) | |
| USER_TEMPLATE = """Create one complete HTML infographic page that follows the generation prompt faithfully and uses the reference image only as visual design guidance. | |
| [Reference Image] | |
| {reference_note} | |
| [Canvas] | |
| {canvas_spec} | |
| [Generation Prompt] | |
| {generation_prompt} | |
| [Output] | |
| Briefly describe your design approach (1-2 sentences), then output the complete HTML in a ```html code block. | |
| """ | |
| CANVAS_CHOICES = [ | |
| "Auto (reference size, else 1280×720)", | |
| "1280×720 (landscape)", | |
| "1536×1024 (landscape)", | |
| "1024×1536 (portrait)", | |
| "1080×1350 (portrait)", | |
| "1024×1024 (square)", | |
| ] | |
| def _orientation(w: int, h: int) -> str: | |
| r = w / h | |
| if r > 1.1: | |
| return "landscape" | |
| if r < 0.9: | |
| return "portrait" | |
| return "square" | |
| def resolve_canvas(canvas: str, image: Image.Image | None) -> tuple[int, int]: | |
| m = re.match(r"(\d+)×(\d+)", canvas or "") | |
| if m: | |
| return int(m.group(1)), int(m.group(2)) | |
| if image is not None: | |
| return image.size | |
| return 1280, 720 | |
| def build_user_text(prompt: str, w: int, h: int, has_image: bool) -> str: | |
| orientation = _orientation(w, h) | |
| canvas_note = f" ({w}x{h}, {orientation})" | |
| reference_note = ( | |
| f"The reference image is attached above{canvas_note}. " | |
| "Use it for visual inspiration only: color palette, layout grid, typographic hierarchy, section arrangement, and overall design quality. " | |
| "Its visible text, numbers, and labels may be inaccurate or incomplete — do not copy them literally unless the generation prompt or inline data contracts require them." | |
| if has_image | |
| else "No reference image is available. Follow the generation prompt's visual direction closely." | |
| ) | |
| canvas_spec = ( | |
| f"Output size: {w}×{h} px ({orientation}). " | |
| f"Your root container element MUST set width: {w}px; height: {h}px; overflow: hidden. " | |
| f"Do not use percentage, vh, vw, or other relative units for the root canvas dimensions." | |
| ) | |
| return USER_TEMPLATE.format( | |
| generation_prompt=prompt.strip(), reference_note=reference_note, canvas_spec=canvas_spec | |
| ) | |
| def extract_html(text: str) -> str | None: | |
| m = re.search(r"```html\s*(.*?)(```|$)", text, flags=re.S | re.I) | |
| body = m.group(1) if m else text | |
| low = body.lower() | |
| i = low.find("<!doctype html") | |
| if i < 0: | |
| i = low.find("<html") | |
| if i < 0: | |
| return None | |
| body = "<!DOCTYPE html>\n" + body[i:] | |
| else: | |
| body = body[i:] | |
| return body.strip() | |
| # --------------------------------------------------------------------------- | |
| # Sandboxed rendering | |
| # --------------------------------------------------------------------------- | |
| CSP = ( | |
| "default-src 'none'; " | |
| "script-src 'unsafe-inline' https://cdn.jsdelivr.net; " | |
| "style-src 'unsafe-inline' https://fonts.googleapis.com https://cdn.jsdelivr.net; " | |
| "font-src data: https://fonts.gstatic.com https://cdn.jsdelivr.net; " | |
| "img-src data: blob: https://placehold.co; " | |
| "connect-src 'none'; form-action 'none';" | |
| ) | |
| def render_iframe(doc: str, w: int, h: int) -> str: | |
| fit = ( | |
| f'<meta http-equiv="Content-Security-Policy" content="{CSP}">' | |
| "<style>html,body{margin:0!important;padding:0!important;overflow:hidden!important;}</style>" | |
| "<script>(function(){function f(){var s=window.innerWidth/" + str(w) + ";" | |
| "document.documentElement.style.zoom=s;}" | |
| "window.addEventListener('resize',f);f();})();</script>" | |
| ) | |
| if re.search(r"<head[^>]*>", doc, flags=re.I): | |
| doc = re.sub(r"(<head[^>]*>)", lambda m: m.group(1) + fit, doc, count=1, flags=re.I) | |
| else: | |
| doc = fit + doc | |
| srcdoc = html_lib.escape(doc, quote=True) | |
| return ( | |
| f'<div style="width:100%;aspect-ratio:{w}/{h};border:1px solid var(--border-color-primary);' | |
| f'border-radius:8px;overflow:hidden;background:#fff">' | |
| f'<iframe sandbox="allow-scripts" referrerpolicy="no-referrer" srcdoc="{srcdoc}" ' | |
| f'style="width:100%;height:100%;border:0;display:block"></iframe></div>' | |
| ) | |
| PLACEHOLDER = ( | |
| '<div style="padding:48px;text-align:center;opacity:.6;border:1px dashed var(--border-color-primary);' | |
| 'border-radius:8px">{}</div>' | |
| ) | |
| def _duration(prompt, reference_image=None, canvas=CANVAS_CHOICES[0], max_new_tokens=6144, *args, **kwargs): | |
| # measured ~30 tok/s on ZeroGPU (5.3k tokens in 181s incl. warm-up) + ~15% margin | |
| return min(300, int((10 + int(max_new_tokens) / 32) * 1.15)) | |
| def generate( | |
| prompt: str, | |
| reference_image: Image.Image | None = None, | |
| canvas: str = CANVAS_CHOICES[0], | |
| max_new_tokens: int = 6144, | |
| temperature: float = 0.8, | |
| top_p: float = 0.95, | |
| seed: int = 0, | |
| ): | |
| """Generate an aesthetic HTML/CSS infographic (slide, poster, dashboard) with AesCode-8B. | |
| Args: | |
| prompt: Detailed content prompt: title, sections, exact text, chart/table data. | |
| reference_image: Optional reference image used as layout/style guidance. | |
| canvas: Canvas size of the root artboard; "Auto" uses the reference image size (or 1280×720). | |
| max_new_tokens: Maximum number of tokens to generate. | |
| temperature: Sampling temperature. | |
| top_p: Nucleus sampling top-p. | |
| seed: Random seed. | |
| Returns: | |
| Rendered preview (HTML iframe), the generated HTML source, a downloadable .html file and the model's design note. | |
| """ | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error("Please enter a content prompt.") | |
| if reference_image is not None and not isinstance(reference_image, Image.Image): | |
| reference_image = Image.open(reference_image) | |
| if reference_image is not None: | |
| reference_image = reference_image.convert("RGB") | |
| w, h = resolve_canvas(canvas, reference_image) | |
| user_text = build_user_text(prompt, w, h, reference_image is not None) | |
| content = [] | |
| if reference_image is not None: | |
| content.append({"type": "image", "image": reference_image}) | |
| content.append({"type": "text", "text": user_text}) | |
| messages = [ | |
| {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}, | |
| {"role": "user", "content": content}, | |
| ] | |
| inputs = processor.apply_chat_template( | |
| messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" | |
| ).to(model.device) | |
| torch.manual_seed(int(seed)) | |
| streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True) | |
| gen_kwargs = dict( | |
| **inputs, | |
| streamer=streamer, | |
| max_new_tokens=int(max_new_tokens), | |
| do_sample=True, | |
| temperature=float(temperature), | |
| top_p=float(top_p), | |
| top_k=0, | |
| ) | |
| thread = Thread(target=model.generate, kwargs=gen_kwargs) | |
| t0 = time.perf_counter() | |
| thread.start() | |
| text = "" | |
| n = 0 | |
| for chunk in streamer: | |
| text += chunk | |
| n += 1 | |
| if n % 40 == 0: | |
| yield ( | |
| gr.update(value=PLACEHOLDER.format(f"Writing HTML… {time.perf_counter() - t0:.0f}s")), | |
| gr.update(value=text), | |
| gr.update(), | |
| gr.update(), | |
| ) | |
| thread.join() | |
| elapsed = time.perf_counter() - t0 | |
| print(f"[aescode] generated {len(text)} chars in {elapsed:.1f}s") | |
| doc = extract_html(text) | |
| m = re.search(r"```html", text, flags=re.I) | |
| note = text[: m.start()].strip() if m else "" | |
| note = (note + f"\n\n*Generated in {elapsed:.0f}s.*").strip() | |
| if doc is None: | |
| yield ( | |
| PLACEHOLDER.format("The model did not return an HTML document — see the raw output."), | |
| text, | |
| None, | |
| note, | |
| ) | |
| return | |
| path = tempfile.NamedTemporaryFile(suffix=".html", delete=False, prefix="aescode_").name | |
| with open(path, "w", encoding="utf-8") as f: | |
| f.write(doc) | |
| yield render_iframe(doc, w, h), doc, path, note | |
| HOSPITAL_PROMPT = open("examples/hospital_dashboard_prompt.txt", encoding="utf-8").read() | |
| EV_IDEA = ( | |
| "A quarterly sales dashboard for a European EV brand. Show revenue and units delivered for the " | |
| "last four quarters, a regional breakdown across Germany, France, Norway and the Netherlands, and a " | |
| "table of the top five models. Keep it clean and corporate, dark text on a light background." | |
| ) | |
| CSS = """ | |
| #col-container { max-width: 1280px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="AesCode-8B") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# AesCode-8B · Aesthetic HTML/CSS infographics\n" | |
| "Describe a slide, poster or dashboard (text, sections, chart & table data) and optionally add a " | |
| "reference image for layout and style. [AesCode-8B](https://huggingface.co/microsoft/AesCode-8B) " | |
| "writes a complete, editable HTML/CSS page (charts via ECharts), rendered below in a sandboxed frame. " | |
| "[Code](https://github.com/microsoft/AesCode)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| prompt = gr.Textbox( | |
| label="Content prompt", | |
| lines=8, | |
| placeholder="Create a landscape dashboard titled \"…\". Section 1: … Render this as an ECharts bar chart with values [ … ].", | |
| ) | |
| with gr.Column(scale=2): | |
| reference_image = gr.Image(label="Reference image (optional)", type="pil", height=230) | |
| with gr.Row(): | |
| canvas = gr.Dropdown(CANVAS_CHOICES, value=CANVAS_CHOICES[0], label="Canvas size", scale=3) | |
| run = gr.Button("Generate", variant="primary", scale=1) | |
| with gr.Accordion("Advanced settings", open=False): | |
| with gr.Row(): | |
| max_new_tokens = gr.Slider(1024, 8192, value=6144, step=256, label="Max new tokens") | |
| temperature = gr.Slider(0.1, 1.2, value=0.8, step=0.05, label="Temperature") | |
| top_p = gr.Slider(0.5, 1.0, value=0.95, step=0.01, label="Top-p") | |
| seed = gr.Number(value=0, precision=0, label="Seed") | |
| preview = gr.HTML(PLACEHOLDER.format("The rendered page will appear here."), label="Rendered page") | |
| note = gr.Markdown() | |
| with gr.Accordion("HTML source", open=False): | |
| code = gr.Code(language="html", label="Generated HTML", max_lines=40) | |
| html_file = gr.File(label="Download .html") | |
| outputs = [preview, code, html_file, note] | |
| run.click( | |
| generate, | |
| inputs=[prompt, reference_image, canvas, max_new_tokens, temperature, top_p, seed], | |
| outputs=outputs, | |
| api_name="generate", | |
| ) | |
| gr.Examples( | |
| label="With reference image (from the AesCode repo)", | |
| examples=[[HOSPITAL_PROMPT, "examples/hospital_dashboard_reference.png"]], | |
| inputs=[prompt, reference_image], | |
| outputs=outputs, | |
| fn=generate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| gr.Examples( | |
| label="Prompt only", | |
| examples=[[EV_IDEA], [HOSPITAL_PROMPT]], | |
| inputs=[prompt], | |
| outputs=outputs, | |
| fn=generate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |