chuchichaestli 0.0.1.dev18__py3-none-any.whl

chuchichaestli/__about__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1.dev18"

chuchichaestli/__init__.py

@@ -0,0 +1,17 @@
+"""This file is part of Chuchichaestli.
+
+Chuchichaestli is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+Chuchichaestli is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with Chuchichaestli.  If not, see <http://www.gnu.org/licenses/>.
+
+Developed by the Intelligent Vision Systems Group at ZHAW.
+"""

chuchichaestli/debug/__init__.py

@@ -0,0 +1,92 @@
+"""Debugging utilities package.
+
+This file is part of Chuchichaestli.
+
+Chuchichaestli is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+Chuchichaestli is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with Chuchichaestli.  If not, see <http://www.gnu.org/licenses/>.
+
+Developed by the Intelligent Vision Systems Group at ZHAW.
+"""
+
+
+def as_bytes(byte_str: str = "MB") -> int:
+    """Convert bytes strings to integers."""
+    match byte_str.lower():
+        case "kib":
+            return 1 << 10
+        case "mib":
+            return 1 << 20
+        case "gib":
+            return 1 << 30
+        case "tib":
+            return 1 << 40
+        case "kb":
+            return 1_000
+        case "mb":
+            return 1_000_000
+        case "gb":
+            return 1_000_000_000
+        case "tb":
+            return 1_000_000_000_000
+        case _:
+            return 1
+
+
+def cli_pbar(
+    r_fill: float,
+    prefix: str | list = "",
+    postfix: str | list = "",
+    bar_length: int = 60,
+    fill_symbol: str = "#",
+    empty_symbol: str = "-",
+    float_fmt: str = "{:.2f}",
+    int_fmt: str = "{:4d}",
+) -> str:
+    """Construct/print a progressbar of given relative length, optionally with labels.
+
+    Args:
+        r_fill (float): rate of progress, a number between 0 and 1
+        prefix (str | list): label before the progressbar
+        postfix (str | list): label after the progressbar
+        bar_length (int): maximum length of the progressbar; automatically determined if negative
+        fill_symbol (str): symbol indicating the progressbar's filled status
+        empty_symbol (str): symbol indicating the progressbar's empty status
+        float_fmt (str): string format for floats in prefix
+        int_fmt (str): string format for ints in prefix
+    """
+    if isinstance(prefix, list | tuple):
+        for i, p in enumerate(prefix):
+            if isinstance(p, float):
+                prefix[i] = float_fmt.format(p)
+            elif isinstance(p, int):
+                prefix[i] = int_fmt.format(p)
+            if not isinstance(prefix[i], str):
+                prefix[i] = f"{prefix[i]}"
+        prefix = " ".join(prefix)
+    if isinstance(postfix, list | tuple):
+        for i, p in enumerate(postfix):
+            if isinstance(p, float):
+                postfix[i] = float_fmt.format(p)
+            elif isinstance(p, int):
+                postfix[i] = int_fmt.format(p)
+            if not isinstance(postfix[i], str):
+                postfix[i] = f"{postfix[i]}"
+        postfix = " ".join(postfix)
+    bar = fill_symbol * int(r_fill * bar_length) + empty_symbol * int(
+        (1 - r_fill) * bar_length
+    )
+    if bar_length > 0:
+        line = f"{prefix} [{bar}] {postfix}"
+    else:
+        line = f"{prefix}\t{postfix}"
+    return line

chuchichaestli/debug/memory_usage.py

@@ -0,0 +1,99 @@
+"""Memory debugging utils.
+
+This file is part of Chuchichaestli.
+
+Chuchichaestli is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+Chuchichaestli is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with Chuchichaestli.  If not, see <http://www.gnu.org/licenses/>.
+
+Developed by the Intelligent Vision Systems Group at ZHAW.
+"""
+
+import torch
+from torch import nn
+from chuchichaestli.debug.visitor import Visitor, Hook, HookDirection
+
+
+class CudaMemoryStatsVisitor(Visitor):
+    """Visitor that collects memory statistics on CUDA GPUs."""
+
+    def __init__(self, **kwargs):
+        """Initialize the visitor."""
+        if "hook_default" in kwargs:
+            hook_default = kwargs.pop("hook_default")
+        else:
+            hook_default = Hook(self._hook, HookDirection.Forward)
+        super().__init__(hook_default=hook_default, **kwargs)
+
+    def _hook(self, module: nn.Module, _input: torch.Tensor, _output: torch.Tensor):
+        """Hook that collects the memory statistics."""
+        mem_dict = torch.cuda.memory_stats()
+        mem_dict["module"] = module._get_name()
+        self._memory_stats.append(mem_dict)
+
+    @staticmethod
+    def has_cuda(verbose: bool = False) -> bool:
+        """Check if CUDA is built and available."""
+        is_avail = torch.cuda.is_available()
+        is_built = torch.backends.cuda.is_built()
+        is_cuda = is_avail and is_built
+        if verbose:
+            if not is_built:
+                if not is_avail:
+                    print(
+                        "CUDA not available most likely because you do not have a CUDA-enabled "
+                        "device on this machine."
+                    )
+                else:
+                    print(
+                        "CUDA not available because the current PyTorch install was not built with "
+                        "CUDA enabled."
+                    )
+        return is_cuda
+
+
+class MPSMemoryAllocationVisitor(Visitor):
+    """Visitor that collects memory statistics on Apple systems."""
+
+    def __init__(self, **kwargs):
+        """Initialize the visitor."""
+        if "hook_default" in kwargs:
+            hook_default = kwargs.pop("hook_default")
+        else:
+            hook_default = Hook(self._hook, HookDirection.Forward)
+        super().__init__(hook_default=hook_default, **kwargs)
+
+    def _hook(self, module: nn.Module, _input: torch.Tensor, _output: torch.Tensor):
+        """Hook that collects the memory statistics."""
+        allocated = torch.mps.current_allocated_memory()
+        mem_dict = {"allocated": allocated, "module": module._get_name()}
+        self._memory_stats.append(mem_dict)
+
+    @staticmethod
+    def has_mps(verbose: bool = False) -> bool:
+        """Check if MPS is built and available."""
+        is_avail = torch.backends.mps.is_available()
+        is_built = torch.backends.mps.is_built()
+        is_mps = is_avail and is_built
+        if verbose:
+            if not is_avail:
+                if not is_built:
+                    print(
+                        "MPS not available because the current PyTorch install was not built with "
+                        "MPS enabled."
+                    )
+                else:
+                    print(
+                        "MPS not available because the current MacOS version is not 12.3+ and/or "
+                        "you do not have an MPS-enabled device on this machine."
+                    )
+        return is_mps

chuchichaestli/debug/visitor.py

@@ -0,0 +1,137 @@
+"""Visitor pattern implementation for PyTorch models.
+
+This file is part of Chuchichaestli.
+
+Chuchichaestli is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+Chuchichaestli is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with Chuchichaestli.  If not, see <http://www.gnu.org/licenses/>.
+
+Developed by the Intelligent Vision Systems Group at ZHAW.
+"""
+
+from dataclasses import dataclass
+from enum import Enum
+
+from collections.abc import Callable
+from typing import Any
+
+from torch import nn
+from chuchichaestli.debug import as_bytes, cli_pbar
+
+
+class HookDirection(Enum):
+    """Enum that represents the direction of a hook."""
+
+    Forward: str = "forward"
+    Backward: str = "backward"
+    Both: str = "both"
+
+
+@dataclass
+class Hook:
+    """Represents a hook (function to call and direction to attach the hook)."""
+
+    fn: Callable[[nn.Module, Any], Any]
+    direction: HookDirection
+
+
+class Visitor:
+    """Generic implementation of a Visitor."""
+
+    def __init__(
+        self,
+        hook_default: Hook | None = None,
+        hook_map: dict[object, Hook] = {},
+        max_depth: int | None = None,
+        _memory_stats: list = [],
+    ):
+        """Initialize the Visitor.
+
+        Args:
+            hook_default: The default hook to use.
+            hook_map: A mapping from layer types to hooks.
+            max_depth: The maximum depth to visit.
+            _memory_stats: Precached memory statistics history.
+        """
+        self.hook_default = hook_default
+        self.hook_map = hook_map
+        self.max_depth = max_depth
+        self._memory_stats = _memory_stats
+
+    @property
+    def memory_stats(self):
+        """Return the collected memory statistics."""
+        return self._memory_stats
+
+    def unlink(self):
+        """Unlink hooks and maps."""
+        self.hook_default = None
+        self.hook_map = {}
+
+    def visit(self, visitee, depth: int = 0, caller: nn.Module | None = None):
+        """Function called when visiting an object."""
+        # print(visitee, depth, caller)
+        if self.max_depth and depth > self.max_depth:
+            return
+        for layer in visitee.modules():
+            if caller is layer:
+                return
+            layer_type = type(layer)
+            hook = self.hook_map.setdefault(layer_type, self.hook_default)
+            match hook.direction:
+                case HookDirection.Forward:
+                    layer.register_forward_hook(hook.fn)
+                case HookDirection.Backward:
+                    layer.register_backward_hook(hook.fn)
+                case HookDirection.Both:
+                    layer.register_forward_hook(hook.fn)
+                    layer.register_backward_hook(hook.fn)
+            self.visit(layer, depth=depth + 1, caller=layer)
+
+    def report(
+        self,
+        unit: str = "MB",
+        with_bar: bool = True,
+        bar_length: int = 60,
+        verbose: bool = True,
+    ) -> list[str] | None:
+        """Report the memory statistics for the visited module(s).
+
+        Args:
+            unit (str): The byte unit; one in [KB, MB, GB, TB, KiB, MiB, GiB, TiB].
+            with_bar (bool): Add simple horizontal percentage bars.
+            bar_length (int): The length of the percentage bars in number of characters.
+            verbose (bool): Directly print lines to stdout.
+        """
+        if not hasattr(self, "memory_stats"):
+            return []
+        divisor = as_bytes(unit)
+        mem_stats = self.memory_stats.copy()
+        if not mem_stats:
+            return []
+        total_alloc = sum(module["allocated"] for module in mem_stats)
+        for stat in mem_stats:
+            stat["rel_alloc"] = stat["allocated"] / total_alloc
+            stat["allocated"] /= divisor
+        total_alloc /= divisor
+
+        lines = []
+        module_name_length = max([len(stat["module"]) for stat in mem_stats])
+
+        for stat in mem_stats:
+            prefix = [stat["module"].ljust(module_name_length), stat["allocated"], unit]
+            postfix = [stat["rel_alloc"] * 100, "%"]
+            line = cli_pbar(stat["rel_alloc"], prefix, postfix, bar_length=bar_length)
+            lines.append(line)
+        if verbose:
+            print("\n".join(lines))
+        return lines

chuchichaestli/diffusion/__init__.py

@@ -0,0 +1,23 @@
+"""Diffusion schedulers sub-module.
+
+This file is part of Chuchichaestli.
+
+Chuchichaestli is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+Chuchichaestli is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with Chuchichaestli.  If not, see <http://www.gnu.org/licenses/>.
+
+Developed by the Intelligent Vision Systems Group at ZHAW.
+"""
+
+from chuchichaestli.diffusion.ddpm import DDPM, InDI, PriorGrad, CFGDDPM, BBDM, DDIM
+
+__all__ = ["DDPM", "InDI", "PriorGrad", "CFGDDPM", "BBDM", "DDIM"]

chuchichaestli/diffusion/ddpm/__init__.py

@@ -0,0 +1,28 @@
+"""Implementation of DDPM and its variants.
+
+This file is part of Chuchichaestli.
+
+Chuchichaestli is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+Chuchichaestli is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with Chuchichaestli.  If not, see <http://www.gnu.org/licenses/>.
+
+Developed by the Intelligent Vision Systems Group at ZHAW.
+"""
+
+from chuchichaestli.diffusion.ddpm.ddpm import DDPM
+from chuchichaestli.diffusion.ddpm.indi import InDI
+from chuchichaestli.diffusion.ddpm.prior_grad import PriorGrad
+from chuchichaestli.diffusion.ddpm.cfg_ddpm import CFGDDPM
+from chuchichaestli.diffusion.ddpm.bbdm import BBDM
+from chuchichaestli.diffusion.ddpm.ddim import DDIM
+
+__all__ = ["DDPM", "InDI", "PriorGrad", "CFGDDPM", "BBDM", "DDIM"]

chuchichaestli/diffusion/ddpm/base.py

@@ -0,0 +1,161 @@
+"""Base class for diffusion processes.
+
+This file is part of Chuchichaestli.
+
+Chuchichaestli is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+Chuchichaestli is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with Chuchichaestli.  If not, see <http://www.gnu.org/licenses/>.
+
+Developed by the Intelligent Vision Systems Group at ZHAW.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+from collections.abc import Generator
+
+import torch
+
+from chuchichaestli.diffusion.distributions import DistributionAdapter
+
+SCHEDULES = {
+    # A simple linear schedule from `beta_start` to `beta_end`
+    "linear": lambda beta_start, beta_end, num_timesteps, device: torch.linspace(
+        beta_start, beta_end, num_timesteps, device=device
+    ),
+    # Scaled linear schedule to account for different step sizes
+    "linear_scaled": lambda beta_start, beta_end, num_timesteps, device: torch.linspace(
+        (1000 / num_timesteps) * beta_start,
+        (1000 / num_timesteps) * beta_end,
+        num_timesteps,
+        device=device,
+    ),
+    # A quadratic growth schedule from `beta_min` to `beta_max`
+    "squared": lambda beta_min, beta_max, num_timesteps, device: beta_min
+    + (
+        (beta_max - beta_min)
+        * (
+            torch.arange(num_timesteps, dtype=torch.float32, device=device)
+            / num_timesteps
+        )
+        ** 2
+    ),
+    # A sigmoid-shaped schedule to achieve a smoother transition from `beta_start` to `beta_end`
+    "sigmoid": lambda beta_start, beta_end, num_timesteps, device: (
+        torch.sigmoid(torch.linspace(-6, 6, num_timesteps, device=device))
+        * (beta_end - beta_start)
+        + beta_start
+    ),
+    # A cosine-shaped schedule for smoother transitions, emphasizing the start and the end
+    "cosine": lambda beta_min, beta_max, num_timesteps, device: beta_min
+    + (
+        (beta_max - beta_min)
+        * (
+            1
+            - torch.cos(
+                (
+                    torch.arange(num_timesteps, dtype=torch.float32, device=device)
+                    / num_timesteps
+                )
+                * torch.pi
+                / 2
+            )
+        )
+    ),
+    # An exponential schedule to smoothly increase the beta value over time
+    "exponential": lambda beta_min, beta_max, num_timesteps, device: beta_min
+    * (
+        (beta_max / beta_min)
+        ** (
+            torch.arange(num_timesteps, dtype=torch.float32, device=device)
+            / num_timesteps
+        )
+    ),
+}
+
+
+SAMPLERS = {
+    "DDIM": "ddim",
+    "default": "default",
+}
+
+
+class DiffusionProcess(ABC):
+    """Base class for diffusion processes."""
+
+    def __init__(
+        self,
+        timesteps: int,
+        device: str = "cpu",
+        noise_distribution: DistributionAdapter | None = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        """Initialize the diffusion process."""
+        self.num_time_steps = timesteps
+        self.device = device
+        self.noise_distribution = noise_distribution
+
+    @abstractmethod
+    def noise_step(
+        self, x_t: torch.Tensor, *args, **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Noise step for the diffusion process.
+
+        Args:
+            x_t: Tensor of shape (batch_size, *).
+            *args: Additional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Tuple of the sampled tensor, noise tensor and timesteps.
+        """
+        pass
+
+    @abstractmethod
+    def generate(
+        self,
+        model: Any,
+        condition: torch.Tensor,
+        n: int = 1,
+        yield_intermediate: bool = False,
+        *args,
+        **kwargs,
+    ) -> Generator[torch.Tensor, None, torch.Tensor]:
+        """Sample from the diffusion process.
+
+        Samples n times the number of conditions from the diffusion process.
+
+        Args:
+            model: Model to use for sampling.
+            condition: Tensor to condition generation on. For unconditional generation, supply a zero-tensor of the sample shape.
+            n: Number of samples to generate (batch size).
+            yield_intermediate: Yield intermediate results. This turns the function into a generator.
+            *args: Additional arguments.
+            **kwargs: Additional keyword
+        """
+        pass
+
+    def sample_timesteps(self, n: int) -> torch.Tensor:
+        """Sample timesteps for the diffusion process."""
+        return torch.randint(0, self.num_time_steps, (n,), device=self.device)
+
+    def sample_noise(self, shape: torch.Size) -> torch.Tensor:
+        """Sample noise for the diffusion process.
+
+        Args:
+            shape: Shape of the noise tensor.
+            mean: Mean of the noise tensor.
+            scale: Scale of the noise tensor.
+        """
+        if self.noise_distribution is not None:
+            return self.noise_distribution(shape)
+        return torch.randn(shape, device=self.device)