Build a Screenshot-to-Code Agent Using LLaVA via Ollama and Open-Source Tools

What We're Building

We're building a local CLI tool that takes a UI screenshot, feeds it to the LLaVA vision model running in Ollama, extracts raw HTML/CSS code from the model's response, and renders a live preview in a browser using Playwright. No API keys. No credit cards. Everything runs on your laptop.

Feature list:

• Single-command pipeline: python run.py --image mockup.png
• LLaVA 13B/7B vision model for zero-shot screenshot-to-code conversion
• Automatic extraction of code blocks from model output (handles markdown fences)
• Playwright-powered headless browser preview that auto-refreshes
• Local file watcher so you can tweak the generated code and see changes instantly
• Full offline capability after the first model pull

Architecture Overview

Here's how the pieces connect. The flow is linear but event-driven: a screenshot enters, a prompt hits LLaVA, the response is sanitized, and Playwright renders the result.

Why this stack? LLaVA is the best open-source vision model that runs comfortably on consumer hardware (16GB RAM for 7B, 32GB for 13B). Ollama handles the GPU offloading and quantization automatically. Playwright gives us a real browser engine without the overhead of Electron or a full IDE.

Prerequisites

All free. All open-source. No cloud dependencies.

Hardware minimum: 16GB RAM for LLaVA 7B. If you have 32GB+, use LLaVA 13B for significantly better layout fidelity.

Step 1: Set Up Ollama and Pull LLaVA

Install Ollama from the official site, then pull the model:

# Install Ollama (macOS/Linux)
curl -fsSL https://ollama.com/install.sh | sh

# For Windows, use the MSI installer from ollama.com

# Pull LLaVA 7B (faster, lower memory)
ollama pull llava:7b

# OR pull LLaVA 13B (better quality)
ollama pull llava:13b

Verify it works with a quick test:

ollama run llava:7b "Describe this image in one sentence." --image ./test-screenshot.png

If you get a coherent description, you're ready. Ollama runs as a background service on localhost:11434.

Step 2: Write the Vision-to-Code Extraction Script

Create extractor.py. This module handles the API call to Ollama and parses the model's output.

import base64
import json
import re
import requests
from pathlib import Path

OLLAMA_URL = "http://localhost:11434/api/generate"

def encode_image(image_path: str) -> str:
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def extract_html_css(raw_text: str) -> str:
    # Try to extract code from markdown fences first
    match = re.search(r"```(?:html)?\n(.*?)```", raw_text, re.DOTALL)
    if match:
        return match.group(1).strip()
    
    # Fallback: assume the entire response is the code
    # Strip any leading/trailing explanatory text
    lines = raw_text.split("\n")
    code_lines = []
    in_code = False
    for line in lines:
        if line.strip().startswith("<"):
            in_code = True
        if in_code:
            code_lines.append(line)
    
    if code_lines:
        return "\n".join(code_lines)
    
    raise ValueError("Could not extract HTML/CSS from model response")

def generate_code_from_screenshot(image_path: str, model: str = "llava:7b") -> str:
    image_b64 = encode_image(image_path)
    
    prompt = """You are an expert frontend developer. Given this UI screenshot, generate the complete HTML and CSS code to recreate it.

Requirements:
- Output ONLY the code, wrapped in a single ```html code block.
- Use semantic HTML5 and modern CSS (flexbox/grid).
- Make it responsive with a max-width container.
- Include all visible text content.
- Match colors, spacing, and typography as closely as possible.
- Do NOT include any explanation before or after the code block."""
    
    payload = {
        "model": model,
        "prompt": prompt,
        "images": [image_b64],
        "stream": False,
        "options": {
            "temperature": 0.1,  # Low temp for deterministic code generation
            "num_predict": 4096
        }
    }
    
    response = requests.post(OLLAMA_URL, json=payload, timeout=120)
    response.raise_for_status()
    result = response.json()
    
    raw_output = result.get("response", "")
    if not raw_output:
        raise RuntimeError("Empty response from LLaVA")
    
    return extract_html_css(raw_output)

Key decisions: Temperature is set to 0.1 because we want deterministic, faithful code—not creative variations. The prompt explicitly requests a code fence, which makes parsing reliable.

Step 3: Build the Live Preview Server with Playwright

Create preview.py. This launches a headless Chromium instance, injects the generated HTML, and optionally watches for file changes.

import asyncio
import os
from pathlib import Path
from playwright.async_api import async_playwright

OUTPUT_FILE = "output.html"
PREVIEW_URL = f"file://{Path(OUTPUT_FILE).resolve()}"

async def launch_preview():
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=False)  # Set to True for CI
        page = await browser.new_page(viewport={"width": 1440, "height": 900})
        
        await page.goto(PREVIEW_URL)
        print(f"Preview launched at {PREVIEW_URL}")
        
        # Watch for file changes and auto-reload
        last_mtime = os.path.getmtime(OUTPUT_FILE)
        try:
            while True:
                await asyncio.sleep(1)
                current_mtime = os.path.getmtime(OUTPUT_FILE)
                if current_mtime != last_mtime:
                    print("Change detected, reloading...")
                    await page.reload()
                    last_mtime = current_mtime
        except KeyboardInterrupt:
            print("Shutting down preview...")
            await browser.close()

def save_html(html_code: str, output_path: str = OUTPUT_FILE):
    with open(output_path, "w") as f:
        f.write(html_code)
    print(f"HTML saved to {output_path}")

Why Playwright over a simple HTTP server? Playwright gives you a real browser viewport, handles CSS rendering exactly like Chrome, and supports auto-reload via file watching. A static HTTP server would require manual refresh and might miss rendering quirks.

Step 4: Orchestrating the Full Pipeline

Create run.py. This ties everything together and adds CLI argument support.

import argparse
import asyncio
import sys
from pathlib import Path

from extractor import generate_code_from_screenshot
from preview import save_html, launch_preview

def main():
    parser = argparse.ArgumentParser(description="Screenshot-to-Code Agent")
    parser.add_argument("--image", "-i", required=True, help="Path to UI screenshot")
    parser.add_argument("--model", "-m", default="llava:7b", 
                        choices=["llava:7b", "llava:13b"],
                        help="LLaVA model variant")
    parser.add_argument("--no-preview", action="store_true", 
                        help="Skip live preview (just generate code)")
    parser.add_argument("--output", "-o", default="output.html",
                        help="Output HTML file path")
    
    args = parser.parse_args()
    
    if not Path(args.image).exists():
        print(f"Error: Image file '{args.image}' not found.")
        sys.exit(1)
    
    print(f"Generating code from {args.image} using {args.model}...")
    html_code = generate_code_from_screenshot(args.image, args.model)
    
    save_html(html_code, args.output)
    
    if not args.no_preview:
        print("Launching live preview...")
        asyncio.run(launch_preview())
    else:
        print(f"Done. Open {args.output} in your browser to view.")

if __name__ == "__main__":
    main()

Running the Agent End-to-End

# Make sure Ollama is running (it usually starts automatically)
ollama serve &

# Take a screenshot of any UI (or use an existing one)
# macOS: Cmd+Shift+4, Windows: Win+Shift+S

# Run the pipeline
python run.py --image ./figma-export.png --model llava:13b

# Output:
# Generating code from ./figma-export.png using llava:13b...
# HTML saved to output.html
# Launching live preview...
# Preview launched at file:///Users/you/project/output.html

A Chromium window opens with your generated UI. Edit output.html in your editor—Playwright detects the file change and reloads instantly. This tight feedback loop is critical for iterating on prompts or cleaning up model output.

Extensions and Production Hardening

1. Batch Processing: Wrap the pipeline in a loop to process an entire folder of screenshots. This is useful for design system migration projects where you have dozens of legacy mockups.

2. Prompt Engineering with Examples: LLaVA supports multi-turn conversations. You can prime it with a few example screenshot-code pairs before the actual request. This dramatically improves fidelity for custom component libraries.

3. Headless CI Integration: Set Playwright headless=True and add this to a GitHub Action. Every time a designer pushes a mockup to a repo, the action generates the HTML and posts it as a PR comment. This is the kind of internal tooling pattern we explore in depth in our guide on debugging in the customer's environment without direct access.

4. Self-Hosted Memory: Want the agent to remember your design system tokens? Combine this pipeline with persistent memory syncing over SSH. The pattern is identical to what we built in Deja Vu: Syncing Agent Memory Over SSH for Persistent, Self-Hosted Coding Workflows.

5. Quality Gate with a Second Model: After LLaVA generates the code, pass it to a smaller text-only model (like Phi-3) that checks for accessibility issues, missing alt text, or semantic HTML violations. This two-pass architecture catches common vision-model blind spots.

Common Pitfalls and Debugging

Model outputs explanations instead of code. Increase the prompt's emphasis on "ONLY the code" and check that your regex handles both html and bare fences. If LLaVA still adds preamble, try the 13B model—it follows instructions more reliably.

Generated layout is completely wrong. Vision models struggle with complex overlapping elements. Simplify the screenshot: crop to a single component or section. For full-page mockups, consider breaking them into header/body/footer chunks and stitching the outputs.

Ollama times out on large images. LLaVA processes images at their native resolution. Resize screenshots to 1024px wide before sending. Add --max-width 1024 to your screenshot tool or use Python's Pillow to resize in the pipeline.

Playwright can't find Chromium. Run playwright install chromium explicitly. On Linux servers, you may need additional system dependencies: playwright install-deps chromium.

CUDA out of memory. If you're on a GPU-poor machine, force CPU inference with OLLAMA_NUM_GPU=0 ollama serve. It's slower but won't crash.

The preview shows a blank page. Check that the extracted HTML has opening <html> and <body> tags. LLaVA sometimes omits them. Add a sanitization step in extract_html_css() that wraps the output in a minimal HTML skeleton if missing.

FAQ

Q: Can I use a different vision model? Yes. Swap llava:7b for llava:13b, bakllava, or any Ollama-compatible vision model. The API payload structure is identical. Just ensure the model supports the images field in /api/generate.

Q: How do I handle interactive elements like dropdowns or modals? LLaVA generates static HTML/CSS from a single screenshot. For interactive states, provide multiple screenshots (e.g., dropdown-open.png, dropdown-closed.png) and ask the model to generate JavaScript that toggles between them.

Q: Is this production-ready? As a prototyping accelerator, absolutely. For production UI code, treat the output as a first draft. It'll get you 70% there on layout and typography. The remaining 30%—accessibility, responsive breakpoints, framework integration—still needs an engineer's touch. This mirrors the enterprise deployment patterns we discuss in our case study on deploying LLM features that survive security review.

Q: Can I fine-tune LLaVA on my company's design system? LLaVA's architecture supports LoRA fine-tuning, but that's a separate workflow involving curated screenshot-code pairs. For teams serious about automating design-to-code, consider building a dataset from your Figma component library exports and your production codebase. The resulting model would generate code that actually matches your internal conventions.

Q: What if I want to schedule this as a recurring job? Wrap run.py in a cron job or a simple polling loop. For more complex scheduling logic—like triggering generation when a new mockup appears in a shared drive—check out the patterns in our calendar-scheduling agent build guide. The orchestration principles transfer directly.