source/app.py

from __future__ import annotations

from pathlib import Path

import gradio as gr
import numpy as np
import torch
from model import TuringSurrogate
from physics import gray_scott_step, initial_state
from PIL import Image
from safetensors.torch import load_file

ARTIFACT = Path(__file__).resolve().parent / "artifacts" / "turing-neural-field"
MODEL = TuringSurrogate()
MODEL.load_state_dict(load_file(ARTIFACT / "model.safetensors"))
MODEL.eval()


def render(field: torch.Tensor) -> np.ndarray:
    values = field[0, 1].numpy()
    normalized = np.clip(values / max(0.05, float(values.max())), 0, 1)
    return np.stack(
        [
            normalized * 70,
            normalized * 210,
            30 + normalized * 225,
        ],
        axis=2,
    ).astype(np.uint8)


def compare(feed: float, kill: float, steps: int, seed: int) -> Image.Image:
    feed_tensor = torch.tensor([feed], dtype=torch.float32)
    kill_tensor = torch.tensor([kill], dtype=torch.float32)
    physics = initial_state(1, 32, seed)
    neural = physics.clone()
    with torch.no_grad():
        for _ in range(steps):
            physics = gray_scott_step(physics, feed_tensor, kill_tensor)
            neural = MODEL(neural, feed_tensor, kill_tensor)
    separator = np.full((32, 2, 3), 245, dtype=np.uint8)
    combined = np.concatenate([render(physics), separator, render(neural)], axis=1)
    return Image.fromarray(combined).resize((1056, 512), Image.Resampling.NEAREST)


with gr.Blocks(title="Turing Neural Field") as demo:
    gr.Markdown("# Turing Neural Field\nLeft: numerical physics. Right: learned surrogate.")
    with gr.Row():
        feed = gr.Slider(0.025, 0.060, value=0.042, step=0.001, label="Feed")
        kill = gr.Slider(0.050, 0.072, value=0.061, step=0.001, label="Kill")
        steps = gr.Slider(1, 160, value=100, step=1, label="Steps")
        seed = gr.Slider(0, 10000, value=4096, step=1, label="Seed")
    output = gr.Image(value=compare(0.042, 0.061, 100, 4096))
    button = gr.Button("Grow pattern", variant="primary")
    button.click(compare, inputs=[feed, kill, steps, seed], outputs=output)


if __name__ == "__main__":
    demo.launch()