Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # noqa: E402 (must precede torch) | |
| import html as html_lib # noqa: E402 | |
| import re # noqa: E402 | |
| import tempfile # noqa: E402 | |
| import time # noqa: E402 | |
| from threading import Thread # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import torch # noqa: E402 | |
| from PIL import Image # noqa: E402 | |
| from transformers import ( # noqa: E402 | |
| AutoProcessor, | |
| Qwen3VLForConditionalGeneration, | |
| TextIteratorStreamer, | |
| ) | |
| MODEL_ID = "microsoft/AesCode-32B" | |
| processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| model = Qwen3VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID, dtype=torch.bfloat16, attn_implementation="sdpa" | |
| ).to("cuda") | |
| model.eval() | |
| # --- Prompting: copied 1:1 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 (match reference image)", | |
| "1280x720 (landscape)", | |
| "1536x1024 (landscape)", | |
| "1024x1536 (portrait)", | |
| "1024x1024 (square)", | |
| ] | |
| def _orientation(w: int, h: int) -> str: | |
| ratio = w / h | |
| if ratio > 1.15: | |
| return "landscape" | |
| if ratio < 0.87: | |
| return "portrait" | |
| return "square" | |
| def _resolve_canvas(canvas: str, image: Image.Image | None): | |
| if canvas.startswith("Auto"): | |
| if image is None: | |
| return None | |
| return image.size | |
| w, h = canvas.split(" ")[0].split("x") | |
| return int(w), int(h) | |
| def build_user_text(prompt: str, has_image: bool, wh) -> str: | |
| if wh: | |
| w, h = wh | |
| orient = _orientation(w, h) | |
| canvas_note = f" ({w}x{h}, {orient})" | |
| canvas_spec = ( | |
| f"Output size: {w}×{h} px ({orient}). " | |
| 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." | |
| ) | |
| else: | |
| canvas_note = "" | |
| canvas_spec = ( | |
| "No specific canvas size given. Default to width: 1280px; height: 720px; overflow: hidden " | |
| "on the root container element." | |
| ) | |
| 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." | |
| ) | |
| 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() | |
| def render_iframe(doc: str | None, wh) -> str: | |
| if not doc: | |
| return "<div class='placeholder'>The rendered page will appear here.</div>" | |
| w, h = wh or (1280, 720) | |
| # Zoom the fixed-size artboard to fit the iframe width. | |
| fit = ( | |
| "<script>(function(){function f(){var z=Math.min(1,window.innerWidth/%d);" | |
| "document.documentElement.style.zoom=z;}f();window.addEventListener('resize',f);})();</script>" | |
| % w | |
| ) | |
| low = doc.lower() | |
| j = low.rfind("</body>") | |
| full = doc[:j] + fit + doc[j:] if j >= 0 else doc + fit | |
| return ( | |
| f"<div class='frame-wrap' style='aspect-ratio:{w}/{h};'>" | |
| f"<iframe sandbox='allow-scripts' srcdoc=\"{html_lib.escape(full, quote=True)}\" " | |
| f"style='width:100%;height:100%;border:0;background:white;'></iframe></div>" | |
| ) | |
| def _duration(prompt, reference_image=None, canvas=CANVAS_CHOICES[0], max_new_tokens=6144, *args, **kwargs): | |
| # Measured ~16 tok/s for bf16 AesCode-32B on a full ZeroGPU card; small margin on top. | |
| return int(20 + int(max_new_tokens) / 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, | |
| ): | |
| """Generate a complete HTML/CSS infographic (slides, posters, dashboards) with AesCode-32B. | |
| Args: | |
| prompt: Detailed content description (titles, sections, chart/table data) of the page to build. | |
| reference_image: Optional reference image used as aesthetic/layout guidance. | |
| canvas: Canvas size; "Auto" matches the reference image (or 1280x720 without one). | |
| max_new_tokens: Maximum number of tokens to generate. | |
| temperature: Sampling temperature. | |
| top_p: Nucleus sampling top-p. | |
| Returns: | |
| Rendered preview (HTML iframe), the generated HTML source, and a downloadable .html file. | |
| """ | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error("Please enter a content prompt.") | |
| if reference_image is not None: | |
| reference_image = reference_image.convert("RGB") | |
| wh = _resolve_canvas(canvas, reference_image) | |
| user_text = build_user_text(prompt, reference_image is not None, wh) | |
| 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("cuda") | |
| 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, | |
| ) | |
| t0 = time.perf_counter() | |
| thread = Thread(target=model.generate, kwargs=gen_kwargs) | |
| thread.start() | |
| text = "" | |
| last = 0.0 | |
| n_chunks = 0 | |
| for chunk in streamer: | |
| text += chunk | |
| n_chunks += 1 | |
| now = time.perf_counter() | |
| if now - last > 0.5: | |
| last = now | |
| yield gr.update(), text, None | |
| thread.join() | |
| elapsed = time.perf_counter() - t0 | |
| n_tok = len(processor.tokenizer(text)["input_ids"]) | |
| print(f"[aescode] generated {n_tok} tokens in {elapsed:.1f}s ({n_tok / max(elapsed, 1e-6):.1f} tok/s)") | |
| doc = extract_html(text) | |
| path = None | |
| if doc: | |
| with tempfile.NamedTemporaryFile("w", suffix=".html", delete=False, encoding="utf-8") as f: | |
| f.write(doc) | |
| path = f.name | |
| else: | |
| gr.Warning("No complete HTML document found in the output (try raising max new tokens).") | |
| yield render_iframe(doc, wh), text, path | |
| HOSPITAL_PROMPT = open("examples/hospital_dashboard_prompt.txt", encoding="utf-8").read() | |
| EV_PROMPT = ( | |
| "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); } | |
| .frame-wrap { width: 100%; border: 1px solid var(--border-color-primary); border-radius: 8px; overflow: hidden; background: white; } | |
| .placeholder { padding: 80px 20px; text-align: center; opacity: .6; border: 1px dashed var(--border-color-primary); border-radius: 8px; } | |
| """ | |
| with gr.Blocks(title="AesCode-32B") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# 🎨 AesCode-32B — aesthetic HTML/CSS infographics\n" | |
| "Describe a slide, poster or dashboard (and optionally attach a reference image for layout & style). " | |
| "[AesCode-32B](https://huggingface.co/microsoft/AesCode-32B) writes a complete, editable HTML/CSS page " | |
| "with real tables and ECharts charts. " | |
| "[Code](https://github.com/microsoft/AesCode) · [8B model](https://huggingface.co/microsoft/AesCode-8B)\n\n" | |
| "Long pages take a few minutes — the source streams in below while the model writes." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| prompt = gr.Textbox( | |
| label="Content prompt", | |
| lines=8, | |
| placeholder="Describe the page: title, sections, the exact numbers for charts and tables, style…", | |
| ) | |
| 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=2) | |
| run = gr.Button("Generate page", variant="primary", scale=1) | |
| with gr.Accordion("Advanced settings", open=False): | |
| 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") | |
| preview = gr.HTML(render_iframe(None, None), label="Rendered page") | |
| with gr.Accordion("Model output (HTML source)", open=False): | |
| code = gr.Code(label="Raw output", language="markdown", lines=20, max_lines=40) | |
| html_file = gr.File(label="Download HTML") | |
| outputs = [preview, code, html_file] | |
| run.click( | |
| generate, | |
| inputs=[prompt, reference_image, canvas, max_new_tokens, temperature, top_p], | |
| outputs=outputs, | |
| api_name="generate", | |
| ) | |
| gr.Examples( | |
| examples=[[HOSPITAL_PROMPT, "examples/hospital_dashboard_reference.png"]], | |
| inputs=[prompt, reference_image], | |
| outputs=outputs, | |
| fn=generate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Example with reference image (from the AesCode repo)", | |
| ) | |
| gr.Examples( | |
| examples=[[EV_PROMPT, "1280x720 (landscape)"]], | |
| inputs=[prompt, canvas], | |
| outputs=outputs, | |
| fn=generate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Text-only example", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |