libcom 0.1.8__py3-none-any.whl

libcom/image_harmonization/source/src/lbm/models/utils.py

@@ -0,0 +1,377 @@
+import logging
+import math
+from copy import deepcopy
+from typing import List, Tuple
+
+import torch
+import torch.nn.functional as F
+
+TILING_METHODS = ["average", "gaussian", "linear"]
+
+
+class Tiler:
+    def get_tiles(
+        self,
+        input: torch.Tensor,
+        tile_size: tuple,
+        overlap_size: tuple,
+        scale: int = 1,
+        out_channels: int = 3,
+    ) -> List[List[torch.tensor]]:
+        """Get tiles
+        Args:
+            input (torch.Tensor): input array of shape (batch_size, channels, height, width)
+            tile_size (tuple): tile size
+            overlap_size (tuple): overlap size
+            scale (int): scaling factor of the output wrt input
+            out_channels (int): number of output channels
+        Returns:
+            List[List[torch.Tensor]]: List of tiles
+        """
+        # assert isinstance(scale, int)
+        assert (
+            overlap_size[0] <= tile_size[0]
+        ), f"Overlap size {overlap_size} must be smaller than tile size {tile_size}"
+        assert (
+            overlap_size[1] <= tile_size[1]
+        ), f"Overlap size {overlap_size} must be smaller than tile size {tile_size}"
+
+        B, C, H, W = input.shape
+        tile_size_H, tile_size_W = tile_size
+
+        # sets overlap to 0 if the input is smaller than the tile size (i.e. no overlap)
+        overlap_H, overlap_W = (
+            overlap_size[0] if H > tile_size_H else 0,
+            overlap_size[1] if W > tile_size_W else 0,
+        )
+
+        self.output_overlap_size = (
+            int(overlap_H * scale),
+            int(overlap_W * scale),
+        )
+        self.tile_size = tile_size
+        self.output_tile_size = (
+            int(tile_size_H * scale),
+            int(tile_size_W * scale),
+        )
+        self.output_shape = (
+            B,
+            out_channels,
+            int(H * scale),
+            int(W * scale),
+        )
+        tiles = []
+        logging.debug(f"(Tiler) Input shape: {(B, C, H, W)}")
+        logging.debug(f"(Tiler) Output shape: {self.output_shape}")
+        logging.debug(f"(Tiler) Tile size: {(tile_size_H, tile_size_W)}")
+        logging.debug(f"(Tiler) Overlap size: {(overlap_H, overlap_W)}")
+        # loop over all tiles in the image with overlap
+        for i in range(0, H, tile_size_H - overlap_H):
+            row = []
+            for j in range(0, W, tile_size_W - overlap_W):
+                tile = deepcopy(
+                    input[
+                        :,
+                        :,
+                        i : i + tile_size_H,
+                        j : j + tile_size_W,
+                    ]
+                )
+                row.append(tile)
+            tiles.append(row)
+        return tiles
+
+    def merge_tiles(
+        self, tiles: List[List[torch.tensor]], tiling_method: str = "gaussian"
+    ) -> torch.tensor:
+        """Merge tiles by averaging the overlaping regions
+        Args:
+            tiles (Dict[str, Tile]): dictionary of processed tiles
+            tiling_method (str): tiling method. Can be "average", "gaussian" or "linear"
+        Returns:
+            torch.tensor: output image
+        """
+        if tiling_method == "average":
+            return self._average_merge_tiles(tiles)
+        elif tiling_method == "gaussian":
+            return self._gaussian_merge_tiles(tiles)
+        elif tiling_method == "linear":
+            return self._linear_merge_tiles(tiles)
+        else:
+            raise ValueError(
+                f"Unknown tiling method {tiling_method}. Available methods are {TILING_METHODS}"
+            )
+
+    def _average_merge_tiles(self, tiles: List[List[torch.tensor]]) -> torch.tensor:
+        """Merge tiles by averaging the overlaping regions
+        Args:
+            tiles (Dict[str, Tile]): dictionary of processed tiles
+        Returns:
+            torch.tensor: output image
+        """
+
+        output = torch.zeros(self.output_shape)
+
+        # weights to store multiplicity
+        weights = torch.zeros(self.output_shape)
+
+        _, _, output_H, output_W = self.output_shape
+        output_overlap_size_H, output_overlap_size_W = self.output_overlap_size
+        output_tile_size_H, output_tile_size_W = self.output_tile_size
+
+        for id_i, i in enumerate(
+            range(
+                0,
+                output_H,
+                output_tile_size_H - output_overlap_size_H,
+            )
+        ):
+            for id_j, j in enumerate(
+                range(
+                    0,
+                    output_W,
+                    output_tile_size_W - output_overlap_size_W,
+                )
+            ):
+                output[
+                    :,
+                    :,
+                    i : i + output_tile_size_H,
+                    j : j + output_tile_size_W,
+                ] += (
+                    tiles[id_i][id_j] * 1
+                )
+                weights[
+                    :,
+                    :,
+                    i : i + output_tile_size_H,
+                    j : j + output_tile_size_W,
+                ] += 1
+
+        # outputs is summed up with this multiplicity
+        # so we need to divide by the weights wich is either 1, 2 or 4 depending on the region
+        output = output / weights
+        return output
+
+    def _gaussian_weights(
+        self, tile_width: int, tile_height: int, nbatches: int, channels: int
+    ):
+        """Generates a gaussian mask of weights for tile contributions.
+
+        Args:
+            tile_width (int): width of the tile
+            tile_height (int): height of the tile
+            nbatches (int): number of batches
+            channels (int): number of channels
+        Returns:
+            torch.tensor: weights
+        """
+        import numpy as np
+        from numpy import exp, pi, sqrt
+
+        latent_width = tile_width
+        latent_height = tile_height
+
+        var = 0.01
+        midpoint = (
+            latent_width - 1
+        ) / 2  # -1 because index goes from 0 to latent_width - 1
+        x_probs = [
+            exp(
+                -(x - midpoint)
+                * (x - midpoint)
+                / (latent_width * latent_width)
+                / (2 * var)
+            )
+            / sqrt(2 * pi * var)
+            for x in range(latent_width)
+        ]
+        midpoint = latent_height / 2
+        y_probs = [
+            exp(
+                -(y - midpoint)
+                * (y - midpoint)
+                / (latent_height * latent_height)
+                / (2 * var)
+            )
+            / sqrt(2 * pi * var)
+            for y in range(latent_height)
+        ]
+
+        weights = np.outer(y_probs, x_probs)
+        return torch.tile(
+            torch.tensor(weights, device="cpu"), (nbatches, channels, 1, 1)
+        )
+
+    def _gaussian_merge_tiles(self, tiles: List[List[torch.tensor]]) -> torch.tensor:
+        """Merge tiles by averaging the overlaping regions
+        Args:
+            List[List[torch.tensor]]: List of processed tiles
+        Returns:
+            torch.tensor: output image
+        """
+        B, output_C, output_H, output_W = self.output_shape
+        output_overlap_size_H, output_overlap_size_W = self.output_overlap_size
+        output_tile_size_H, output_tile_size_W = self.output_tile_size
+
+        output = torch.zeros(self.output_shape)
+        # weights to store multiplicity
+        weights = torch.zeros(self.output_shape)
+
+        for id_i, i in enumerate(
+            range(
+                0,
+                output_H,
+                output_tile_size_H - output_overlap_size_H,
+            )
+        ):
+            for id_j, j in enumerate(
+                range(
+                    0,
+                    output_W,
+                    output_tile_size_W - output_overlap_size_W,
+                )
+            ):
+                w = self._gaussian_weights(
+                    tiles[id_i][id_j].shape[3],
+                    tiles[id_i][id_j].shape[2],
+                    B,
+                    output_C,
+                )
+                output[
+                    :,
+                    :,
+                    i : i + output_tile_size_H,
+                    j : j + output_tile_size_W,
+                ] += (
+                    tiles[id_i][id_j] * w
+                )
+                weights[
+                    :,
+                    :,
+                    i : i + output_tile_size_H,
+                    j : j + output_tile_size_W,
+                ] += w
+
+        # outputs is summed up with this multiplicity
+        output = output / weights
+        return output
+
+    def _blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[
+                :, :, y, :
+            ] * (y / blend_extent)
+        return b
+
+    def _blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[
+                :, :, :, x
+            ] * (x / blend_extent)
+        return b
+
+    def _linear_merge_tiles(self, tiles: List[List[torch.tensor]]) -> torch.Tensor:
+        """Merge tiles by blending the overlaping regions
+        Args:
+            tiles (List[List[torch.tensor]]): List of processed tiles
+        Returns:
+            torch.Tensor: output image
+        """
+        output_overlap_size_H, output_overlap_size_W = self.output_overlap_size
+        output_tile_size_H, output_tile_size_W = self.output_tile_size
+
+        res_rows = []
+        tiles_copy = deepcopy(tiles)
+
+        # Cut the right and bottom overlap region
+        limit_i = output_tile_size_H - output_overlap_size_H
+        limit_j = output_tile_size_W - output_overlap_size_W
+        for i, tile_row in enumerate(tiles_copy):
+            res_row = []
+            for j, tile in enumerate(tile_row):
+                tile_val = tile
+                if j > 0:
+                    tile_val = self._blend_h(
+                        tile_row[j - 1], tile, output_overlap_size_W
+                    )
+                tiles_copy[i][j] = tile_val
+                if i > 0:
+                    tile_val = self._blend_v(
+                        tiles_copy[i - 1][j], tile_val, output_overlap_size_H
+                    )
+                tiles_copy[i][j] = tile_val
+                res_row.append(tile_val[:, :, :limit_i, :limit_j])
+            res_rows.append(torch.cat(res_row, dim=3))
+        output = torch.cat(res_rows, dim=2)
+        return output
+
+
+def extract_into_tensor(
+    a: torch.Tensor, t: torch.Tensor, x_shape: Tuple[int, ...]
+) -> torch.Tensor:
+    """
+    Extracts values from a tensor into a new tensor using indices from another tensor.
+
+    :param a: the tensor to extract values from.
+    :param t: the tensor containing the indices.
+    :param x_shape: the shape of the tensor to extract values into.
+    :return: a new tensor containing the extracted values.
+    """
+
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def pad(x: torch.Tensor, base_h: int, base_w: int) -> torch.Tensor:
+    """
+    Pads a tensor to the nearest multiple of base_h and base_w.
+
+    :param x: the tensor to pad.
+    :param base_h: the base height.
+    :param base_w: the base width.
+    :return: the padded tensor.
+    """
+    h, w = x.shape[-2:]
+    h_ = math.ceil(h / base_h) * base_h
+    w_ = math.ceil(w / base_w) * base_w
+    if w_ != w:
+        x = F.pad(x, (0, abs(w_ - w), 0, 0))
+    if h_ != h:
+        x = F.pad(x, (0, 0, 0, abs(h_ - h)))
+    return x
+
+
+def append_dims(x: torch.Tensor, target_dims: int) -> torch.Tensor:
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(
+            f"input has {x.ndim} dims but target_dims is {target_dims}, which is less"
+        )
+    return x[(...,) + (None,) * dims_to_append]
+
+
+@torch.no_grad()
+def update_ema(
+    target_params: List[torch.Tensor],
+    source_params: List[torch.Tensor],
+    rate: float = 0.99,
+):
+    """
+    Update target parameters to be closer to those of source parameters using
+    an exponential moving average.
+
+    :param target_params: the target parameter sequence.
+    :param source_params: the source parameter sequence.
+    :param rate: the EMA rate (closer to 1 means slower).
+    """
+    for targ, src in zip(target_params, source_params):
+        targ.detach().mul_(rate).add_(src, alpha=1 - rate)

libcom/image_harmonization/source/src/lbm/models/vae/__init__.py

@@ -0,0 +1,4 @@
+from .autoencoderKL import AutoencoderKLDiffusers
+from .autoencoderKL_config import AutoencoderKLDiffusersConfig
+
+__all__ = ["AutoencoderKLDiffusers", "AutoencoderKLDiffusersConfig"]

libcom/image_harmonization/source/src/lbm/models/vae/autoencoderKL.py

@@ -0,0 +1,136 @@
+import torch
+from diffusers.models import AutoencoderKL
+
+from ..base.base_model import BaseModel
+from ..utils import Tiler, pad
+from .autoencoderKL_config import AutoencoderKLDiffusersConfig
+
+
+class AutoencoderKLDiffusers(BaseModel):
+    """This is the VAE class used to work with latent models
+
+    Args:
+
+        config (AutoencoderKLDiffusersConfig): The config class which defines all the required parameters.
+    """
+
+    def __init__(self, config: AutoencoderKLDiffusersConfig):
+        BaseModel.__init__(self, config)
+        self.config = config
+        self.vae_model = AutoencoderKL.from_pretrained(
+            config.version,
+            subfolder=config.subfolder,
+            revision=config.revision,
+        )
+        self.tiling_size = config.tiling_size
+        self.tiling_overlap = config.tiling_overlap
+
+        # get downsampling factor
+        self._get_properties()
+
+    @torch.no_grad()
+    def _get_properties(self):
+        self.has_shift_factor = (
+            hasattr(self.vae_model.config, "shift_factor")
+            and self.vae_model.config.shift_factor is not None
+        )
+        self.shift_factor = (
+            self.vae_model.config.shift_factor if self.has_shift_factor else 0
+        )
+
+        # set latent channels
+        self.latent_channels = self.vae_model.config.latent_channels
+        self.has_latents_mean = (
+            hasattr(self.vae_model.config, "latents_mean")
+            and self.vae_model.config.latents_mean is not None
+        )
+        self.has_latents_std = (
+            hasattr(self.vae_model.config, "latents_std")
+            and self.vae_model.config.latents_std is not None
+        )
+        self.latents_mean = self.vae_model.config.latents_mean
+        self.latents_std = self.vae_model.config.latents_std
+
+        x = torch.randn(1, self.vae_model.config.in_channels, 32, 32)
+        z = self.encode(x)
+
+        # set downsampling factor
+        self.downsampling_factor = int(x.shape[2] / z.shape[2])
+
+    def encode(self, x: torch.tensor, batch_size: int = 8):
+        latents = []
+        for i in range(0, x.shape[0], batch_size):
+            latents.append(
+                self.vae_model.encode(x[i : i + batch_size]).latent_dist.sample()
+            )
+        latents = torch.cat(latents, dim=0)
+        latents = (latents - self.shift_factor) * self.vae_model.config.scaling_factor
+
+        return latents
+
+    def decode(self, z: torch.tensor):
+
+        if self.has_latents_mean and self.has_latents_std:
+            latents_mean = (
+                torch.tensor(self.latents_mean)
+                .view(1, self.latent_channels, 1, 1)
+                .to(z.device, z.dtype)
+            )
+            latents_std = (
+                torch.tensor(self.latents_std)
+                .view(1, self.latent_channels, 1, 1)
+                .to(z.device, z.dtype)
+            )
+            z = z * latents_std / self.vae_model.config.scaling_factor + latents_mean
+        else:
+            z = z / self.vae_model.config.scaling_factor + self.shift_factor
+
+        use_tiling = (
+            z.shape[2] > self.tiling_size[0] or z.shape[3] > self.tiling_size[1]
+        )
+
+        if use_tiling:
+            samples = []
+            for i in range(z.shape[0]):
+
+                z_i = z[i].unsqueeze(0)
+
+                tiler = Tiler()
+                tiles = tiler.get_tiles(
+                    input=z_i,
+                    tile_size=self.tiling_size,
+                    overlap_size=self.tiling_overlap,
+                    scale=self.downsampling_factor,
+                    out_channels=3,
+                )
+
+                for i, tile_row in enumerate(tiles):
+                    for j, tile in enumerate(tile_row):
+                        tile_shape = tile.shape
+                        # pad tile to inference size if tile is smaller than inference size
+                        tile = pad(
+                            tile,
+                            base_h=self.tiling_size[0],
+                            base_w=self.tiling_size[1],
+                        )
+                        tile_decoded = self.vae_model.decode(tile).sample
+                        tiles[i][j] = (
+                            tile_decoded[
+                                0,
+                                :,
+                                : int(tile_shape[2] * self.downsampling_factor),
+                                : int(tile_shape[3] * self.downsampling_factor),
+                            ]
+                            .cpu()
+                            .unsqueeze(0)
+                        )
+
+                # merge tiles
+                samples.append(tiler.merge_tiles(tiles=tiles))
+
+            samples = torch.cat(samples, dim=0)
+
+        else:
+            samples = self.vae_model.decode(z).sample
+
+        return samples

libcom/image_harmonization/source/src/lbm/models/vae/autoencoderKL_config.py

@@ -0,0 +1,27 @@
+from typing import Tuple
+
+from pydantic.dataclasses import dataclass
+
+from ..base import ModelConfig
+
+
+@dataclass
+class AutoencoderKLDiffusersConfig(ModelConfig):
+    """This is the VAEConfig class which defines all the useful parameters to instantiate the model.
+
+    Args:
+
+        version (str): The version of the model. Defaults to "stabilityai/sdxl-vae".
+        subfolder (str): The subfolder of the model if loaded from another model. Defaults to "".
+        revision (str): The revision of the model. Defaults to "main".
+        input_key (str): The key of the input data in the batch. Defaults to "image".
+        tiling_size (Tuple[int, int]): The size of the tiling. Defaults to (64, 64).
+        tiling_overlap (Tuple[int, int]): The overlap of the tiling. Defaults to (16, 16).
+    """
+
+    version: str = "stabilityai/sdxl-vae"
+    subfolder: str = ""
+    revision: str = "main"
+    input_key: str = "image"
+    tiling_size: Tuple[int, int] = (64, 64)
+    tiling_overlap: Tuple[int, int] = (16, 16)

libcom/image_harmonization/source/src/lbm/trainer/__init__.py

@@ -0,0 +1,85 @@
+"""
+This module contains the training pipeline and the training configuration along with all relevant parts 
+of the training pipeline such as loggers and callbacks.
+
+The :mod:`cr.trainer` includes the following submodules:
+
+- :mod:`cr.trainer.trainer`: the main training pipeline class for ClipDrop.
+- :mod:`cr.trainer.training_config`: the configuration for the training pipeline.
+- :mod:`cr.trainer.loggers`: the loggers for logging samples to wandb.
+
+
+Examples
+########
+
+Train a model using the training pipeline
+
+.. code-block:: python
+
+    from cr.trainer import TrainingPipeline, TrainingConfig
+    from cr.data import DataPipeline, DataConfig
+    from pytorch_lightning import Trainer
+    from cr.data.datasets import DataModule, DataModuleConfig
+
+    # Create a model to train
+    model = DummyModel()
+
+    # Create a training configuration
+    config = TrainingConfig(
+        experiment_id="test",
+        optimizers_name=["AdamW"],
+        optimizers_kwargs=[{}],
+        learning_rates=[1e-3],
+        lr_schedulers_name=[None],
+        lr_schedulers_kwargs=[{}],
+        trainable_params=[["./*"]],
+        log_keys="txt",
+        log_samples_model_kwargs={
+            "max_samples": 8,
+            "num_steps": 20,
+            "input_shape": (4, 32, 32),
+            "guidance_scale": 7.5,
+        }
+    )
+
+    # Create a training pipeline
+    pipeline = TrainingPipeline(model=model, pipeline_config=config)
+    
+    # Create a DataModule
+    data_module = DataModule(
+            train_config=DataModuleConfig(
+                shards_path_or_urls="your urls or paths",
+                decoder="pil",
+                shuffle_buffer_size=100,
+                per_worker_batch_size=32,
+                num_workers=4,
+            ),
+            train_filters_mappers=your_mappers_and_filters,
+            eval_config=DataModuleConfig(
+                shards_path_or_urls="your urls or paths",
+                decoder="pil",
+                shuffle_buffer_size=100,
+                per_worker_batch_size=32,
+                num_workers=4,
+            ),
+            eval_filters_mappers=your_mappers_and_filters,
+        )
+
+    # Create a trainer
+    trainer = Trainer(
+            accelerator="cuda",
+            max_epochs=1,
+            devices=1,
+            log_every_n_steps=1,
+            default_root_dir="your dir",
+            max_steps=2,
+        )
+
+    # Train the model
+    trainer.fit(pipeline, data_module)
+"""
+
+from .trainer import TrainingPipeline
+from .training_config import TrainingConfig
+
+__all__ = ["TrainingPipeline", "TrainingConfig"]