File size: 2,163 Bytes
626ab93 | 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 | 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()
|