feat: added requirements and sampling

latent_space_explorer/sampling.py

@@ -0,0 +1,69 @@
+from abc import abstractmethod
+
+import torch
+import numpy as np
+
+from .utils import recursive_find_device, recursive_find_dtype
+
+class EncodingSampler:
+    """
+    Class to sample encodings given low dimensional spatial relationships.
+    """
+    def __init__(self, encodes):
+        self.encodes = encodes
+
+    def apply_coefs(self, coefs):
+        """
+        Linear combination of encodings given coefs
+        """
+        device = recursive_find_device(self.encodes)
+        dtype = recursive_find_dtype(self.encodes)
+        # NOTE: Convert from float64 first to `dtype` and *then* to `device` to
+        # prevent issues with certain devices not supporting f64
+        # (*cough cough* Apple)
+        coefs = torch.from_numpy(coefs).to(dtype).to(device)
+
+        def single_apply(encodes):
+            if encodes is None:
+                return None
+            elif len(encodes.shape) == 3:
+                return (coefs[:,None,None] * encodes).sum(0)
+            elif len(encodes.shape) == 2:
+                return (coefs[:,None] * encodes).sum(0)
+            else:
+                raise ValueError("Encoding Sampler couldn't figure out shape of encodings")
+            
+        if isinstance(self.encodes, list) or isinstance(self.encodes, tuple):
+            return list(map(single_apply, self.encodes))
+        else:
+            return single_apply(self.encodes)
+
+    @abstractmethod
+    def __call__(self, point, other_points):
+        """
+        :param point: Point in low space representing user input ([2,] array)
+        :param other_points: Points in low space representing existing prompts ([N,2] array)
+        """
+        pass
+
+class DistanceSampling(EncodingSampler):
+    """
+    Sample based on distances between points in low dim space
+    """
+    def __call__(self, point, other_points):
+        coefs = 1. / ((1. + np.linalg.norm(point[None,:] - other_points, axis = 1) ** 2))
+        return self.apply_coefs(coefs)
+    
+class CircleSampling(EncodingSampler):
+    """
+    Sampler that views all encodings as points on a unit circle
+    """
+    def __call__(self, point, other_points):
+        # Idea: weight of points in same direction should be 1
+        # weight of points in opposite should be 0
+        cos_sims = point @ other_points.transpose() # [2] x [2, N] -> N
+
+        # Negative values don't work, but we want something analagous for "negative signals"
+        # tanh is like -x for low values, but then caps out at 1
+        #cos_sims = np.where(cos_sims<0, np.tanh(cos_sims), cos_sims)
+        return self.apply_coefs(cos_sims)

latent_space_explorer/utils.py

@@ -0,0 +1,41 @@
+import random
+import math
+import torch
+
+def random_circle_init(min_r : float = 0.5, on_edge : bool = False):
+    theta = random.uniform(0, 2 * math.pi)
+    if on_edge:
+        r = 1.0
+    else:
+        r = random.uniform(min_r, 1.0)
+    x = r * math.cos(theta)
+    y = r * math.sin(theta)
+
+    return x, y
+
+def recursive_find_dtype(x):
+    """
+    Assuming x is some list/tuple of things that could be tensors, searches for any tensors and returns dtype
+    """
+    for i in x:
+        if isinstance(i, list):
+            res = recursive_find_dtype(i)
+            if res is None:
+                continue
+            else:
+                return res
+        elif isinstance(i, torch.Tensor):
+            return i.dtype
+
+def recursive_find_device(x):
+    """
+    Assuming x is some list/tuple of things that could be tensors, searches for any tensors and returns device
+    """
+    for i in x:
+        if isinstance(i, list):
+            res = recursive_find_device(i)
+            if res is None:
+                continue
+            return res
+        elif isinstance(i, torch.Tensor):
+            return i.device

main.py

@@ -0,0 +1,20 @@
+from latent_space_explorer import GameConfig, LatentSpaceExplorer
+
+if __name__ == "__main__":
+    config = GameConfig(
+        call_every = 100
+    )
+
+    explorer = LatentSpaceExplorer(config)
+
+    explorer.set_prompts(
+        [
+            "A photo of a cat",
+            "A space-aged ferrari",
+            "artwork of the titanic hitting an iceberg",
+            "a photo of a dog"
+        ]
+    )
+
+    while True:
+        explorer.update()

requirements.txt

@@ -0,0 +1,4 @@
+diffusers
+pygame
+torch
+torchvision