import gradio as gr import os, glob from functools import partial import glob import torch from torch import nn from PIL import Image import numpy as np device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') class RuleCA(nn.Module): def __init__(self, hidden_n=6, rule_channels=4, zero_w2=True, device=device): super().__init__() # The hard-coded filters: self.filters = torch.stack([torch.tensor([[0.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,0.0]]), torch.tensor([[-1.0,0.0,1.0],[-2.0,0.0,2.0],[-1.0,0.0,1.0]]), torch.tensor([[-1.0,0.0,1.0],[-2.0,0.0,2.0],[-1.0,0.0,1.0]]).T, torch.tensor([[1.0,2.0,1.0],[2.0,-12,2.0],[1.0,2.0,1.0]])]).to(device) self.chn = 4 self.rule_channels = rule_channels self.w1 = nn.Conv2d(4*4+rule_channels, hidden_n, 1).to(device) self.relu = nn.ReLU() self.w2 = nn.Conv2d(hidden_n, 4, 1, bias=False).to(device) if zero_w2: self.w2.weight.data.zero_() self.device = device def perchannel_conv(self, x, filters): '''filters: [filter_n, h, w]''' b, ch, h, w = x.shape y = x.reshape(b*ch, 1, h, w) y = torch.nn.functional.pad(y, [1, 1, 1, 1], 'circular') y = torch.nn.functional.conv2d(y, filters[:,None]) return y.reshape(b, -1, h, w) def forward(self, x, rule=0, update_rate=0.5): b, ch, xsz, ysz = x.shape rule_grid = torch.zeros(b, self.rule_channels, xsz, ysz).to(self.device) rule_grid[:,rule] = 1 y = self.perchannel_conv(x, self.filters) # Apply the filters y = torch.cat([y, rule_grid], dim=1) y = self.w2(self.relu(self.w1(y))) # pass the result through out 'brain' b, c, h, w = y.shape update_mask = (torch.rand(b, 1, h, w).to(self.device)+update_rate).floor() return x+y*update_mask def forward_w_rule_grid(self, x, rule_grid, update_rate=0.5): y = self.perchannel_conv(x, self.filters) # Apply the filters y = torch.cat([y, rule_grid], dim=1) y = self.w2(self.relu(self.w1(y))) # pass the result through out 'brain' b, c, h, w = y.shape update_mask = (torch.rand(b, 1, h, w).to(self.device)+update_rate).floor() return x+y*update_mask def to_rgb(self, x): # TODO: rename this to_rgb & explain return x[...,:3,:,:]+0.5 def seed(self, n, sz=128): """Initializes n 'grids', size sz. In this case all 0s.""" return torch.zeros(n, self.chn, sz, sz).to(self.device) def to_frames(video_file): os.system('rm -r guide_frames;mkdir guide_frames') os.system(f"ffmpeg -i {video_file} guide_frames/%04d.jpg") def update(preset, enhance, scale2x, video_file): # Load presets ca = RuleCA(hidden_n=32, rule_channels=3) ca_fn = '' if preset == 'Glowing Crystals': ca_fn = 'glowing_crystals.pt' elif preset == 'Rainbow Diamonds': ca_fn = 'rainbow_diamonds.pt' elif preset == 'Dark Diamonds': ca_fn = 'dark_diamonds.pt' elif preset == 'Dragon Scales': ca = RuleCA(hidden_n=16, rule_channels=3) ca_fn = 'dragon_scales.pt' ca.load_state_dict(torch.load(ca_fn, map_location=device)) # Get video frames to_frames(video_file) size=(426, 240) vid_size = Image.open(f'guide_frames/0001.jpg').size if vid_size[0]>vid_size[1]: # Change < to > if larger side should be capped at 256px size = (256, int(256*(vid_size[1]/vid_size[0]))) else: size = (int(256*(vid_size[0]/vid_size[1])), 256) if scale2x: size = (size[0]*2, size[1]*2) # Starting grid x = torch.zeros(1, 4, size[1], size[0]).to(ca.device) os.system("rm -r steps;mkdir steps") for i in range(2*len(glob.glob('guide_frames/*.jpg'))-1): # load frame im = Image.open(f'guide_frames/{i//2+1:04}.jpg').resize(size) # make rule grid rule_grid = torch.tensor(np.array(im)/255).permute(2, 0, 1).unsqueeze(0).to(ca.device) if enhance: rule_grid = rule_grid * 2 - 0.3 # Add * 2 - 0.3 to 'enhance' an effect # Apply the updates with torch.no_grad(): x = ca.forward_w_rule_grid(x, rule_grid.float()) if i%2==0: img = ca.to_rgb(x).detach().cpu().clip(0, 1).squeeze().permute(1, 2, 0) img = Image.fromarray(np.array(img*255).astype(np.uint8)) img.save(f'steps/{i//2:05}.jpeg') # Write output video from saved frames os.system("ffmpeg -y -v 0 -framerate 24 -i steps/%05d.jpeg video.mp4") return 'video.mp4' demo = gr.Blocks() with demo: gr.Markdown("Choose a preset below, upload a video and then click **Run** to see the output. Read [this report](https://wandb.ai/johnowhitaker/nca/reports/Fun-with-Neural-Cellular-Automata--VmlldzoyMDQ5Mjg0) for background on this project, or check out my [AI art course](https://github.com/johnowhitaker/aiaiart) for an in-depth lesson on Neural Cellular Automata like this.") with gr.Row(): preset = gr.Dropdown(['Glowing Crystals', 'Rainbow Diamonds', 'Dark Diamonds', 'Dragon Scales'], label='Preset') with gr.Column(): enhance = gr.Checkbox(label='Rescale inputs (more extreme results)') scale2x = gr.Checkbox(label='Larger output (slower)') with gr.Row(): inp = gr.Video(format='mp4', source='upload', label="Input video (ideally <30s)") out = gr.Video(label="Output") btn = gr.Button("Run") btn.click(fn=update, inputs=[preset, enhance, scale2x, inp], outputs=out) demo.launch(enable_queue=True)