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 document with all CSS in a single inline "
""
)
if re.search(r"]*>", doc, flags=re.I):
doc = re.sub(r"(]*>)", 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''
f'
'
)
PLACEHOLDER = (
'{}
'
)
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))
@spaces.GPU(duration=_duration)
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)