Spaces:
Build error
Build error
| import random | |
| import torch | |
| import math | |
| def rotate_axis(x, add_angle=0, axis=1): # TODO Replace with a rotation matrix # But this is more fun | |
| axes = list(range(3)) | |
| axes.remove(axis) | |
| ax1, ax2 = axes | |
| angle = torch.atan2(x[..., ax1], x[..., ax2]) | |
| if isinstance(add_angle, torch.Tensor): | |
| while add_angle.ndim < angle.ndim: | |
| add_angle = add_angle.unsqueeze(-1) | |
| angle = angle + add_angle | |
| dist = x.norm(dim=-1) | |
| t = [] | |
| _, t = zip(*sorted([ | |
| (axis, x[..., axis]), | |
| (ax1, torch.sin(angle) * dist), | |
| (ax2, torch.cos(angle) * dist), | |
| ])) | |
| return torch.stack(t, dim=-1) | |
| noise_level = 0.5 | |
| # stolen from https://gist.github.com/ac1b097753f217c5c11bc2ff396e0a57 | |
| # ported from https://github.com/pvigier/perlin-numpy/blob/master/perlin2d.py | |
| def rand_perlin_2d(shape, res, fade=lambda t: 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3): | |
| delta = (res[0] / shape[0], res[1] / shape[1]) | |
| d = (shape[0] // res[0], shape[1] // res[1]) | |
| grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim=-1) % 1 | |
| angles = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1) | |
| gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim=-1) | |
| tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0], | |
| 0).repeat_interleave( | |
| d[1], 1) | |
| dot = lambda grad, shift: ( | |
| torch.stack((grid[:shape[0], :shape[1], 0] + shift[0], grid[:shape[0], :shape[1], 1] + shift[1]), | |
| dim=-1) * grad[:shape[0], :shape[1]]).sum(dim=-1) | |
| n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0]) | |
| n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0]) | |
| n01 = dot(tile_grads([0, -1], [1, None]), [0, -1]) | |
| n11 = dot(tile_grads([1, None], [1, None]), [-1, -1]) | |
| t = fade(grid[:shape[0], :shape[1]]) | |
| return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1]) | |
| def rand_perlin_2d_octaves(shape, res, octaves=1, persistence=0.5): | |
| noise = torch.zeros(shape) | |
| frequency = 1 | |
| amplitude = 1 | |
| for _ in range(octaves): | |
| noise += amplitude * rand_perlin_2d(shape, (frequency * res[0], frequency * res[1])) | |
| frequency *= 2 | |
| amplitude *= persistence | |
| noise *= random.random() - noise_level # haha | |
| noise += random.random() - noise_level # haha x2 | |
| return noise | |
| def load_clip(model_name="ViT-B/16", device="cuda:0" if torch.cuda.is_available() else "cpu"): | |
| import clip | |
| model, preprocess = clip.load(model_name, device=device, jit=False) | |
| if len(preprocess.transforms) > 4: | |
| preprocess.transforms = preprocess.transforms[-1:] | |
| return model, preprocess | |