359 lines
9.0 KiB
Python
359 lines
9.0 KiB
Python
import inspect
|
|
from abc import ABC, abstractmethod
|
|
from collections import deque
|
|
from dataclasses import dataclass, field
|
|
from typing import (Annotated, Any, Callable, Iterable, Self, Sequence, get_args,
|
|
get_origin)
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from lark import Lark, Transformer
|
|
|
|
SCALING_FACTOR = 0.18215
|
|
|
|
|
|
@dataclass
|
|
class LatentMap:
|
|
"""List of loaded latents."""
|
|
|
|
#: Shape of the tensors in the map
|
|
shape: tuple[int, ...]
|
|
|
|
#: Default tensor name
|
|
default: str
|
|
|
|
#: Tensor mapping
|
|
map: dict[str, torch.Tensor] = field(default_factory=dict)
|
|
|
|
@classmethod
|
|
def create(cls, args: Iterable[torch.Tensor]) -> Self:
|
|
tensors = {}
|
|
default: torch.Tensor | None = None
|
|
for idx, arg in enumerate(args):
|
|
name = f"l{idx+1}"
|
|
if default is None:
|
|
default = arg
|
|
tensors[name] = arg
|
|
|
|
else:
|
|
H, W = default.shape[-2:]
|
|
tensors[name] = F.interpolate(
|
|
arg, size=(H, W), mode="bilinear", align_corners=True
|
|
)
|
|
|
|
assert default is not None
|
|
return cls(tuple(default.shape), "l1", tensors)
|
|
|
|
|
|
class NoResolve:
|
|
"""Annotation instructing the Operation() class not to resolve the argument."""
|
|
|
|
#: Global instance
|
|
_instance: Self | None = None
|
|
|
|
def __new__(cls) -> Self:
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
|
|
return cls._instance
|
|
|
|
|
|
class AST(ABC):
|
|
@abstractmethod
|
|
def __call__(self, lm: LatentMap) -> torch.Tensor:
|
|
pass
|
|
|
|
|
|
class Operation(AST):
|
|
def __init__(
|
|
self, name: str, func: Callable[..., torch.Tensor], *args: Any
|
|
) -> None:
|
|
self._name = name
|
|
self._func = func
|
|
self._args = tuple(args)
|
|
|
|
def __str__(self) -> str:
|
|
args = list(str(v) for v in self._args)
|
|
return f"{self._name}({', '.join(args)})"
|
|
|
|
def __repr__(self) -> str:
|
|
args = list(repr(v) for v in self._args)
|
|
return f"{self._name.capitalize()}({', '.join(args)})"
|
|
|
|
def __call__(self, lm: LatentMap) -> torch.Tensor:
|
|
args = []
|
|
remain = deque(self._args)
|
|
|
|
for param in inspect.signature(self._func).parameters.values():
|
|
dont_resolve = False
|
|
dont_resolve = (
|
|
(orig := get_origin(param.annotation)) is not None
|
|
and orig == Annotated
|
|
and (NoResolve() in get_args(param.annotation) or ())
|
|
)
|
|
latent_map = param.annotation == LatentMap
|
|
if not dont_resolve and latent_map:
|
|
resolved = lm
|
|
|
|
else:
|
|
next = remain.popleft()
|
|
if dont_resolve:
|
|
assert isinstance(next, str)
|
|
resolved = next
|
|
|
|
elif isinstance(next, str):
|
|
resolved = lm.map[next]
|
|
|
|
elif isinstance(next, float):
|
|
if param.annotation == torch.Tensor:
|
|
resolved = torch.full(lm.shape, next)
|
|
|
|
else:
|
|
resolved = next
|
|
|
|
elif isinstance(next, Operation):
|
|
resolved = next(lm)
|
|
|
|
else:
|
|
resolved = next
|
|
|
|
assert param.kind in {param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD}
|
|
args.append(resolved)
|
|
|
|
assert not remain
|
|
return self._func(*args)
|
|
|
|
|
|
class Script(AST):
|
|
def __init__(self, ops: Iterable[AST]) -> None:
|
|
self._ops = tuple(ops)
|
|
|
|
def __str__(self) -> str:
|
|
return ";\n".join(str(op) for op in self._ops) + ";"
|
|
|
|
def __repr__(self) -> str:
|
|
return f"{type(self).__name__}({self._ops!r})"
|
|
|
|
def __call__(self, lm: LatentMap) -> torch.Tensor:
|
|
last = self._ops[0](lm)
|
|
for op in self._ops[1:]:
|
|
last = op(lm)
|
|
|
|
return last
|
|
|
|
|
|
operations: dict[str, Callable[..., Operation]] = {}
|
|
|
|
|
|
def operation(func: Callable[..., torch.Tensor], /) -> Callable[..., Operation]:
|
|
"""Create an operations function."""
|
|
|
|
def _wrapper(*args) -> Operation:
|
|
return Operation(func.__name__, func, *args)
|
|
|
|
global operations
|
|
operations[func.__name__] = _wrapper
|
|
return _wrapper
|
|
|
|
|
|
@operation
|
|
def assign(
|
|
lm: LatentMap, out: Annotated[str, NoResolve()], expr: torch.Tensor
|
|
) -> torch.Tensor:
|
|
lm.map[out] = expr
|
|
return expr
|
|
|
|
|
|
@operation
|
|
def mul(l: torch.Tensor, r: torch.Tensor) -> torch.Tensor:
|
|
return l * r
|
|
|
|
|
|
@operation
|
|
def div(l: torch.Tensor, r: torch.Tensor) -> torch.Tensor:
|
|
return l / r
|
|
|
|
|
|
@operation
|
|
def paren(l: torch.Tensor) -> torch.Tensor:
|
|
return l
|
|
|
|
|
|
@operation
|
|
def neg(l: torch.Tensor) -> torch.Tensor:
|
|
return -l
|
|
|
|
|
|
@operation
|
|
def add(l: torch.Tensor, r: torch.Tensor) -> torch.Tensor:
|
|
return l + r
|
|
|
|
|
|
@operation
|
|
def sub(l: torch.Tensor, r: torch.Tensor) -> torch.Tensor:
|
|
return l - r
|
|
|
|
|
|
@operation
|
|
def noise(lm: LatentMap) -> torch.Tensor:
|
|
return torch.randn(lm.shape) * SCALING_FACTOR
|
|
|
|
|
|
@operation
|
|
def chanmap(l: torch.Tensor, map: Sequence[int | None]) -> torch.Tensor:
|
|
output = torch.zeros_like(l)
|
|
for target, source in enumerate(map):
|
|
if source is not None:
|
|
output[:, target, :, :] = l[:, source, :, :]
|
|
|
|
return output
|
|
|
|
|
|
@operation
|
|
def clamp(l: torch.Tensor, min: float, max: float) -> torch.Tensor:
|
|
return torch.clamp(l, min, max)
|
|
|
|
|
|
@operation
|
|
def roll(l: torch.Tensor, shifts: int, dim: int) -> torch.Tensor:
|
|
return torch.roll(l, shifts=shifts, dims=dim)
|
|
|
|
|
|
@operation
|
|
def shuffle(l: torch.Tensor) -> torch.Tensor:
|
|
perm = torch.randperm(l.shape[1])
|
|
return l[:, perm, :, :]
|
|
|
|
|
|
@operation
|
|
def flip(l: torch.Tensor, dim: int) -> torch.Tensor:
|
|
return torch.flip(l, dims=[dim])
|
|
|
|
|
|
@operation
|
|
def sigmoid(l: torch.Tensor) -> torch.Tensor:
|
|
return torch.sigmoid(l)
|
|
|
|
|
|
@operation
|
|
def relu(l: torch.Tensor) -> torch.Tensor:
|
|
return torch.relu(l)
|
|
|
|
|
|
@operation
|
|
def zoom(lm: LatentMap, l: torch.Tensor, scale: float) -> torch.Tensor:
|
|
H, W = l.shape[:-2]
|
|
new_H = int(H * scale)
|
|
new_W = int(W * scale)
|
|
key0 = None
|
|
for key in lm.map.keys():
|
|
if key0 is not None:
|
|
key0 = key
|
|
|
|
lm.map[key] = F.interpolate(
|
|
lm.map[key], size=(new_H, new_W), mode="bilinear", align_corners=True
|
|
)
|
|
|
|
assert key0 is not None
|
|
return lm.map[key0]
|
|
|
|
|
|
@operation
|
|
def resize(lm: LatentMap, height: int, width: int) -> torch.Tensor:
|
|
key0 = None
|
|
for key in lm.map.keys():
|
|
if key0 is not None:
|
|
key0 = key
|
|
|
|
lm.map[key] = F.interpolate(
|
|
l,
|
|
size=(height, width),
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
|
|
assert key0 is not None
|
|
return lm.map[key0]
|
|
|
|
|
|
_parser = Lark("""
|
|
start: ( assignment ";" )* expr
|
|
|
|
assignment: NAME "=" expr -> assign
|
|
|
|
expr: product
|
|
| expr "+" product -> add
|
|
| expr "-" product -> sub
|
|
|
|
product: atom
|
|
| product "*" atom -> mul
|
|
| product "/" atom -> div
|
|
|
|
atom: NUMBER
|
|
| NAME
|
|
| "(" expr ")" -> paren
|
|
| "-" atom -> neg
|
|
| function
|
|
|
|
function: "noise" "(" ")" -> noise
|
|
| "chanmap" "(" expr "," ( OPTCHAN+ ) ")" -> chanmap
|
|
| "clamp" "(" expr "," NUMBER "," NUMBER ")" -> clamp
|
|
| "roll" "(" expr "," INT "," INT ")" -> roll
|
|
| "shuffle" "(" expr ")" -> shuffle
|
|
| "flip" "(" expr "," INT "," INT ")" -> flip
|
|
| "sigmoid" "(" expr ")" -> sigmoid
|
|
| "relu" "(" expr ")" -> relu
|
|
| "resize" "(" NUMBER ")" -> zoom
|
|
| "resize" "(" INT "," INT ")" -> resize
|
|
|
|
NAME: /[a-zA-Z_][a-zA-Z0-9_]*/
|
|
OPTCHAN: CHANNEL | NULL
|
|
CHANNEL: "0" .. "9" | "A" .. "Z"
|
|
NULL: "-"
|
|
|
|
%import common.SIGNED_NUMBER -> NUMBER
|
|
%import common.SIGNED_INT -> INT
|
|
%import common.WS
|
|
%ignore WS
|
|
""")
|
|
|
|
|
|
class _Transformer(Transformer):
|
|
"""AST transformer."""
|
|
|
|
def __default__(self, data, children, meta):
|
|
global operations
|
|
if data in operations:
|
|
return operations[data](*children)
|
|
|
|
return children[0]
|
|
|
|
def start(self, terms) -> Script:
|
|
return Script(terms)
|
|
|
|
def chanmap(self, items) -> AST:
|
|
return chanmap(items[0], items[1:])
|
|
|
|
def NAME(self, token) -> str:
|
|
return str(token)
|
|
|
|
def INT(self, token) -> int:
|
|
return int(token)
|
|
|
|
def NUMBER(self, token) -> float:
|
|
return float(token)
|
|
|
|
def OPTCHAN(self, token) -> int | None:
|
|
if token == "-":
|
|
return None
|
|
|
|
try:
|
|
return int(token)
|
|
|
|
except ValueError:
|
|
return 10 + (ord(token) - ord('A'))
|
|
|
|
|
|
def parse(s: str, /) -> AST:
|
|
return _Transformer().transform(_parser.parse(s))
|