File size: 3,514 Bytes
10d1c5c 6a59e82 537bc18 b20c7d2 537bc18 10d1c5c 537bc18 10d1c5c 537bc18 10d1c5c 537bc18 10d1c5c 537bc18 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import html
import re
import gradio as gr
import spaces
import torch
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
MODEL_ID = "google/diffusiongemma-26B-A4B-it"
if not torch.cuda.is_available():
raise RuntimeError("CUDA is not available — model would silently load on CPU.")
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = DiffusionGemmaForBlockDiffusion.from_pretrained(
MODEL_ID,
dtype="auto",
device_map="auto",
)
PROMPT_TEMPLATE = (
"Eres un desarrollador front-end experto. Genera una página HTML completa y "
"autocontenida en un solo archivo (CSS y JavaScript inline, sin recursos "
"externos) según la siguiente descripción. Responde ÚNICAMENTE con el código "
"HTML, sin explicaciones ni bloques de markdown.\n\n"
"Descripción: {description}"
)
def extract_html(text):
# Keep only the last model turn and drop the thought channel.
if "<|turn>model\n" in text:
text = text.rsplit("<|turn>model\n", 1)[-1]
if "<channel|>" in text:
text = text.rsplit("<channel|>", 1)[-1]
for tok in ("<eos>", "<pad>", "<turn|>", "<bos>"):
text = text.replace(tok, "")
text = text.strip()
# Unwrap a markdown code fence if the model added one anyway.
fenced = re.search(r"```(?:html)?\s*(.*?)```", text, re.DOTALL)
if fenced:
text = fenced.group(1).strip()
return text
@spaces.GPU(duration=120)
def generate(description, max_new_tokens):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": PROMPT_TEMPLATE.format(description=description)}
],
}
]
inputs = processor.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
).to(model.device)
with torch.inference_mode():
output = model.generate(**inputs, max_new_tokens=max_new_tokens)
decoded = processor.decode(output[0], skip_special_tokens=False)
if isinstance(decoded, list):
decoded = decoded[0]
code = extract_html(decoded)
preview = (
f'<iframe srcdoc="{html.escape(code, quote=True)}" sandbox="allow-scripts" '
'style="width:100%;height:70vh;border:1px solid #ccc;border-radius:8px;'
'background:white"></iframe>'
)
return preview, code
with gr.Blocks(title="DiffusionGemma HTML Generator") as demo:
gr.Markdown(
"# DiffusionGemma → HTML\n"
"Describe lo que quieres y [google/diffusiongemma-26B-A4B-it]"
"(https://huggingface.co/google/diffusiongemma-26B-A4B-it) genera la página. "
"El resultado se renderiza abajo en un iframe aislado."
)
with gr.Row():
description = gr.Textbox(
label="¿Qué quieres construir?",
placeholder="ej: una landing page para una cafetería, con menú y formulario de contacto",
lines=2,
scale=4,
)
btn = gr.Button("Generar", variant="primary", scale=1)
max_tokens = gr.Slider(
minimum=256, maximum=4096, value=2048, step=256, label="Max new tokens"
)
with gr.Tab("Vista previa"):
preview = gr.HTML()
with gr.Tab("Código"):
code = gr.Code(language="html")
btn.click(generate, [description, max_tokens], [preview, code])
description.submit(generate, [description, max_tokens], [preview, code])
if __name__ == "__main__":
demo.launch()
|