--- license: gemma base_model: google/diffusiongemma-26B-A4B-it pipeline_tag: text-generation library_name: transformers tags: - diffusion-language-model - generative-ui - openui - openui-lang - gemma language: - en --- # OUI-1 **OUI-1** is the first diffusion model built for generative UI. It is a finetune of Google's DiffusionGemma 26B-A4B-it that writes user interface screens in [openui-lang](https://openui.com/docs), the declarative UI language behind [OpenUI](https://github.com/thesysdev/openui). It scores 71.7% on [Generative UI Benchmark](https://www.openui.com/benchmarks), 5.5x its base model, with 4B active parameters. Given a component library's signatures in the system prompt and a plain-language brief, it returns the screen as code, one component per line, wired into a root. Any component library with signatures can be the prompt, and any OpenUI app can point at it. Generative UI is unusually sensitive to latency, and this is a text diffusion model: it writes a 256-token block at once, starting from noise and committing each token the moment it is sure of it, so a screen arrives in about a second on one GPU. | | | |---|---| | Base model | google/diffusiongemma-26B-A4B-it (26B total, 4B active) | | Method | LoRA finetuning, merged into the base weights; bf16 safetensors | | Context | 16,384 tokens as served | | License | [Gemma Terms of Use](https://ai.google.dev/gemma/terms). This is a Gemma derivative; the base license and its use restrictions apply. | | Publisher | Thesys | ## Results on the Generative UI Benchmark The [generative-ui-bench](https://github.com/thesysdev/generative-ui-bench) protocol: 46 screen briefs in five size bands, 4 generations each, thinking off, one shared system prompt for every model, scored by the benchmark's own validator. Base DiffusionGemma 24/184 (13.0%); OUI-1 132/184 (71.7%). Raw outputs for OUI-1 are committed to the benchmark repo so the number can be rescored offline. ![OpenUI score vs active parameters, open-weight models up to 31B active](scoreboard.png) Every number here was measured with the serving settings below (vLLM 0.24, FP8, the checkpoint's own sampler at 48 denoising steps). One request at a time on an A100 80GB, a light screen takes about a second and a dense one three to six seconds, prompt included. ## How to use The model reads the component library from the system prompt and answers with an openui-lang program. Build the system prompt for your own component library with `@openuidev/cli` (`npx @openuidev/cli generate --out system-prompt.txt`), or use the benchmark's reference prompt (`protocols/openui/prompt.ts` in generative-ui-bench) to try it as is. Render the output with `@openuidev/react-lang` (or the Vue and Svelte renderers) and validate it with `@openuidev/lang-core`. ### vLLM (recommended) vLLM 0.24 or newer. ```bash pip install "vllm>=0.24" vllm serve thesysdev/OUI-1 --trust-remote-code --max-model-len 16384 --quantization fp8 \ --served-model-name OUI-1 --max-num-seqs 4 \ --enable-auto-tool-choice --tool-call-parser gemma4 ``` The sampler settings come from the checkpoint: the 256-token canvas from `config.json`, the entropy-bound sampler (entropy bound 0.1) and 48 denoising steps from `generation_config.json`. No other flags or environment variables are needed; every number above was measured with this line. The weights take 25.8 GiB at FP8; vLLM then fills the rest of the card with KV cache by default, so pass `--gpu-memory-utilization` if the model has to share a GPU. On Ampere (A100) there is no native FP8, so vLLM uses weight-only FP8 through Marlin: the memory saving is real, the compute speedup is not. Then call it as an OpenAI-compatible chat model. `temperature` and `seed` are ignored: the sampler runs the checkpoint's own schedule and no per-request seed is plumbed through, so two identical requests can return differently worded screens. The benchmark runs at `max_tokens` 8192; 4096 covers every screen in the bands above. ```python from openai import OpenAI client = OpenAI(base_url="http://localhost:8000/v1", api_key="local") system = open("system-prompt.txt").read() # your component library, from @openuidev/cli brief = """Status page for the platform team. Build a single screen for this. It must show: 1. current uptime percentage for the API this month 2. a short note on the most recent incident and when it was resolved Cover every numbered item.""" r = client.chat.completions.create( model="OUI-1", messages=[{"role": "system", "content": system}, {"role": "user", "content": brief}], max_tokens=4096, stream=False, ) print(r.choices[0].message.content) # openui-lang: one component per line, wired into root ``` `stream=True` works, with one difference from an autoregressive model: text arrives one 256-token canvas at a time, so a screen that fits in one canvas arrives as a single chunk. To change the step cap use `--diffusion-config '{"canvas_length":256,"max_denoising_steps":32}'` (32 and 48 score the same on the benchmark, 16 breaks the reference graph). `--hf-overrides` keys such as `diffusion_max_denoising_steps` are not read; vLLM takes the step count from `--diffusion-config` or, failing that, from `generation_config.json`. ### Function calling The model keeps Gemma 4's native tool-call format. With `--tool-call-parser gemma4` it returns standard OpenAI `tool_calls` for tools you pass in the request, then writes the screen from the tool result on the next turn. It needs one instruction in the system prompt to prefer the tool over inventing data, e.g.: "When the user asks about the weather or a stock price you MUST call the matching tool first and output nothing else in that turn; write the screen from the returned values on the next turn." In our test (weather, stock price tools) it called the right tool with the right arguments on 9 of 9 asks, never called one on 3 unrelated asks, and used the returned numbers verbatim. Leave `tool_choice` at its default, `"auto"`. `"required"` and a named function are accepted but ignored (vLLM has no structured outputs for diffusion models yet): the model answers with an ordinary screen, and with `"required"` the response carries `finish_reason: "tool_calls"` while `tool_calls` is empty, so do not branch on `finish_reason` alone. ```python import json TOOL_RULE = ("\n\nWhen the user asks about the weather or a stock price you MUST call the " "matching tool first and output nothing else in that turn; write the screen " "from the returned values on the next turn.") tools = [{"type": "function", "function": {"name": "get_weather", "description": "Current weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}] msgs = [{"role": "system", "content": system + TOOL_RULE}, {"role": "user", "content": "weather in Goa this weekend"}] r = client.chat.completions.create(model="OUI-1", tools=tools, max_tokens=4096, messages=msgs) m = r.choices[0].message call = m.tool_calls[0] # get_weather({"city": "Goa"}) result = {"city": "Goa", "condition": "Partly cloudy", "high_c": 31, "low_c": 25, "rain_chance_pct": 20} msgs.append({"role": "assistant", "content": m.content or "", "tool_calls": [{"id": call.id, "type": "function", "function": {"name": call.function.name, "arguments": call.function.arguments}}]}) msgs.append({"role": "tool", "tool_call_id": call.id, "content": json.dumps(result)}) r2 = client.chat.completions.create(model="OUI-1", tools=tools, max_tokens=4096, messages=msgs) print(r2.choices[0].message.content) # the screen, built from the tool result ``` ### Transformers Transformers 5.11 or newer. `AutoModelForCausalLM` does not resolve this architecture, so import the model class directly. bf16 peaks at about 52 GiB of GPU memory (one A100 80GB or H100). `generate` chains 256-token canvases and stops at EOS; keep `max_new_tokens` above 256 or the screen is cut off after the first canvas. ```python import torch from transformers import AutoProcessor from transformers.models.diffusion_gemma.modeling_diffusion_gemma import DiffusionGemmaForBlockDiffusion model_id = "thesysdev/OUI-1" processor = AutoProcessor.from_pretrained(model_id) model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="cuda") messages = [{"role": "system", "content": open("system-prompt.txt").read()}, {"role": "user", "content": brief}] ids = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda") gc = model.generation_config # entropy-bound sampler, 48 steps, from the checkpoint gc.max_new_tokens = 256 * 32 # up to 32 canvases; EOS stops earlier out = model.generate(input_ids=ids, generation_config=gc) print(processor.decode(out.sequences[0, ids.shape[1]:], skip_special_tokens=True)) ``` ### Adapter The LoRA adapter that produced these weights is in `adapter/` (PEFT format, r=64, alpha=128, base `unsloth/diffusiongemma-26B-A4B-it`, a mirror of `google/diffusiongemma-26B-A4B-it`). It is a tied LoRA: its 205 modules target the decoder projections, and DiffusionGemma's encoder shares the same weight storage, so the merge adapts both passes. Plain `PeftModel.from_pretrained(base, adapter)` wraps the decoder only and gives a different model; call `.merge_and_unload()` to reproduce these weights, or use the merged checkpoint above. ## Intended use Generating UI screens for OpenUI-based applications where latency matters and a small, self-hosted model is preferred. Not a general chat model. Outputs should be validated by the openui-lang parser and rendered through a component library that checks its props.