libcom 0.1.8__py3-none-any.whl

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

@@ -0,0 +1,222 @@
+import logging
+import os
+from typing import List, Optional
+
+import torch
+import yaml
+from diffusers import FlowMatchEulerDiscreteScheduler
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+from lbm.models.embedders import (
+    ConditionerWrapper,
+    LatentsConcatEmbedder,
+    LatentsConcatEmbedderConfig,
+)
+from lbm.models.lbm import LBMConfig, LBMModel
+from lbm.models.unets import DiffusersUNet2DCondWrapper
+from lbm.models.vae import AutoencoderKLDiffusers, AutoencoderKLDiffusersConfig
+
+
+def get_model(
+    model_dir: str,
+    save_dir: Optional[str] = None,
+    torch_dtype: torch.dtype = torch.bfloat16,
+    device: str = "cuda",
+) -> LBMModel:
+    """Download the model from the model directory using either a local path or a path to HuggingFace Hub
+
+    Args:
+        model_dir (str): The path to the model directory containing the model weights and config, can be a local path or a path to HuggingFace Hub
+        save_dir (Optional[str]): The local path to save the model if downloading from HuggingFace Hub. Defaults to None.
+        torch_dtype (torch.dtype): The torch dtype to use for the model. Defaults to torch.bfloat16.
+        device (str): The device to use for the model. Defaults to "cuda".
+
+    Returns:
+        LBMModel: The loaded model
+    """
+    if not os.path.exists(model_dir):
+        local_dir = snapshot_download(
+            model_dir,
+            local_dir=save_dir,
+        )
+        model_dir = local_dir
+
+    model_files = os.listdir(model_dir)
+
+    # check yaml config file is present
+    yaml_file = [f for f in model_files if f.endswith(".yaml")]
+    if len(yaml_file) == 0:
+        raise ValueError("No yaml file found in the model directory.")
+
+    # check safetensors weights file is present
+    safetensors_files = sorted([f for f in model_files if f.endswith(".safetensors")])
+    ckpt_files = sorted([f for f in model_files if f.endswith(".ckpt")])
+    if len(safetensors_files) == 0 and len(ckpt_files) == 0:
+        raise ValueError("No safetensors or ckpt file found in the model directory")
+
+    if len(model_files) == 0:
+        raise ValueError("No model files found in the model directory")
+
+    with open(os.path.join(model_dir, yaml_file[0]), "r") as f:
+        config = yaml.safe_load(f)
+
+    model = _get_model_from_config(**config, torch_dtype=torch_dtype)
+
+    if len(safetensors_files) > 0:
+        logging.info(f"Loading safetensors file: {safetensors_files[-1]}")
+        sd = load_file(os.path.join(model_dir, safetensors_files[-1]))
+        model.load_state_dict(sd, strict=True)
+    elif len(ckpt_files) > 0:
+        logging.info(f"Loading ckpt file: {ckpt_files[-1]}")
+        sd = torch.load(
+            os.path.join(model_dir, ckpt_files[-1]),
+            map_location="cpu",
+        )["state_dict"]
+        sd = {k[6:]: v for k, v in sd.items() if k.startswith("model.")}
+        model.load_state_dict(
+            sd,
+            strict=True,
+        )
+    model.to(device).to(torch_dtype)
+
+    model.eval()
+
+    return model
+
+
+def _get_model_from_config(
+    backbone_signature: str = "stabilityai/stable-diffusion-xl-base-1.0",
+    vae_num_channels: int = 4,
+    unet_input_channels: int = 4,
+    timestep_sampling: str = "log_normal",
+    selected_timesteps: Optional[List[float]] = None,
+    prob: Optional[List[float]] = None,
+    conditioning_images_keys: Optional[List[str]] = [],
+    conditioning_masks_keys: Optional[List[str]] = [],
+    source_key: str = "source_image",
+    target_key: str = "source_image_paste",
+    bridge_noise_sigma: float = 0.0,
+    logit_mean: float = 0.0,
+    logit_std: float = 1.0,
+    pixel_loss_type: str = "lpips",
+    latent_loss_type: str = "l2",
+    latent_loss_weight: float = 1.0,
+    pixel_loss_weight: float = 0.0,
+    torch_dtype: torch.dtype = torch.bfloat16,
+    **kwargs,
+):
+
+    conditioners = []
+
+    denoiser = DiffusersUNet2DCondWrapper(
+        in_channels=unet_input_channels,  # Add downsampled_image
+        out_channels=vae_num_channels,
+        center_input_sample=False,
+        flip_sin_to_cos=True,
+        freq_shift=0,
+        down_block_types=[
+            "DownBlock2D",
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+        ],
+        mid_block_type="UNetMidBlock2DCrossAttn",
+        up_block_types=["CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"],
+        only_cross_attention=False,
+        block_out_channels=[320, 640, 1280],
+        layers_per_block=2,
+        downsample_padding=1,
+        mid_block_scale_factor=1,
+        dropout=0.0,
+        act_fn="silu",
+        norm_num_groups=32,
+        norm_eps=1e-05,
+        cross_attention_dim=[320, 640, 1280],
+        transformer_layers_per_block=[1, 2, 10],
+        reverse_transformer_layers_per_block=None,
+        encoder_hid_dim=None,
+        encoder_hid_dim_type=None,
+        attention_head_dim=[5, 10, 20],
+        num_attention_heads=None,
+        dual_cross_attention=False,
+        use_linear_projection=True,
+        class_embed_type=None,
+        addition_embed_type=None,
+        addition_time_embed_dim=None,
+        num_class_embeds=None,
+        upcast_attention=None,
+        resnet_time_scale_shift="default",
+        resnet_skip_time_act=False,
+        resnet_out_scale_factor=1.0,
+        time_embedding_type="positional",
+        time_embedding_dim=None,
+        time_embedding_act_fn=None,
+        timestep_post_act=None,
+        time_cond_proj_dim=None,
+        conv_in_kernel=3,
+        conv_out_kernel=3,
+        projection_class_embeddings_input_dim=None,
+        attention_type="default",
+        class_embeddings_concat=False,
+        mid_block_only_cross_attention=None,
+        cross_attention_norm=None,
+        addition_embed_type_num_heads=64,
+    ).to(torch_dtype)
+
+    if conditioning_images_keys != [] or conditioning_masks_keys != []:
+
+        latents_concat_embedder_config = LatentsConcatEmbedderConfig(
+            image_keys=conditioning_images_keys,
+            mask_keys=conditioning_masks_keys,
+        )
+        latent_concat_embedder = LatentsConcatEmbedder(latents_concat_embedder_config)
+        latent_concat_embedder.freeze()
+        conditioners.append(latent_concat_embedder)
+
+        # Wrap conditioners and set to device
+    conditioner = ConditionerWrapper(
+        conditioners=conditioners,
+    )
+
+    ## VAE ##
+    # Get VAE model
+    vae_config = AutoencoderKLDiffusersConfig(
+        version=backbone_signature,
+        subfolder="vae",
+        tiling_size=(128, 128),
+    )
+    vae = AutoencoderKLDiffusers(vae_config).to(torch_dtype)
+    vae.freeze()
+    vae.to(torch_dtype)
+
+    ## Diffusion Model ##
+    # Get diffusion model
+    config = LBMConfig(
+        source_key=source_key,
+        target_key=target_key,
+        latent_loss_weight=latent_loss_weight,
+        latent_loss_type=latent_loss_type,
+        pixel_loss_type=pixel_loss_type,
+        pixel_loss_weight=pixel_loss_weight,
+        timestep_sampling=timestep_sampling,
+        logit_mean=logit_mean,
+        logit_std=logit_std,
+        selected_timesteps=selected_timesteps,
+        prob=prob,
+        bridge_noise_sigma=bridge_noise_sigma,
+    )
+
+    sampling_noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        backbone_signature,
+        subfolder="scheduler",
+    )
+
+    model = LBMModel(
+        config,
+        denoiser=denoiser,
+        sampling_noise_scheduler=sampling_noise_scheduler,
+        vae=vae,
+        conditioner=conditioner,
+    ).to(torch_dtype)
+
+    return model

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

@@ -0,0 +1,4 @@
+from .base_model import BaseModel
+from .model_config import ModelConfig
+
+__all__ = ["BaseModel", "ModelConfig"]

libcom/image_harmonization/source/src/lbm/models/base/base_model.py

@@ -0,0 +1,66 @@
+from typing import Any, Dict
+
+import torch
+import torch.nn as nn
+
+from .model_config import ModelConfig
+
+
+class BaseModel(nn.Module):
+    def __init__(self, config: ModelConfig):
+        nn.Module.__init__(self)
+        self.config = config
+        self.input_key = config.input_key
+        self.device = torch.device("cpu")
+        self.dtype = torch.float32
+
+    def on_fit_start(self, device: torch.device | None = None, *args, **kwargs):
+        """Called when the training starts
+
+        Args:
+            device (Optional[torch.device], optional): The device to use. Usefull to set
+                relevant parameters on the model and embedder to the right device only
+                once at the start of the training. Defaults to None.
+        """
+        if device is not None:
+            self.device = device
+        self.to(self.device)
+
+    def forward(self, batch: Dict[str, Any], *args, **kwargs):
+        raise NotImplementedError("forward method is not implemented")
+
+    def freeze(self):
+        """Freeze the model"""
+        self.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
+        self = super().to(
+            device=device,
+            dtype=dtype,
+            non_blocking=non_blocking,
+        )
+
+        if device is not None:
+            self.device = device
+        if dtype is not None:
+            self.dtype = dtype
+        return self
+
+    def compute_metrics(self, batch: Dict[str, Any], *args, **kwargs):
+        """Compute the metrics"""
+        return {}
+
+    def sample(self, batch: Dict[str, Any], *args, **kwargs):
+        """Sample from the model"""
+        return {}
+
+    def log_samples(self, batch: Dict[str, Any], *args, **kwargs):
+        """Log the samples"""
+        return None
+
+    def on_train_batch_end(self, batch: Dict[str, Any], *args, **kwargs):
+        """Update the model an optimization is perforned on a batch."""
+        pass

libcom/image_harmonization/source/src/lbm/models/base/model_config.py

@@ -0,0 +1,8 @@
+from pydantic.dataclasses import dataclass
+
+from ...config import BaseConfig
+
+
+@dataclass
+class ModelConfig(BaseConfig):
+    input_key: str = "image"

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

@@ -0,0 +1,4 @@
+from .conditioners_wrapper import ConditionerWrapper
+from .latents_concat import LatentsConcatEmbedder, LatentsConcatEmbedderConfig
+
+__all__ = ["LatentsConcatEmbedder", "LatentsConcatEmbedderConfig", "ConditionerWrapper"]

libcom/image_harmonization/source/src/lbm/models/embedders/base/__init__.py

@@ -0,0 +1,4 @@
+from .base_conditioner import BaseConditioner
+from .base_conditioner_config import BaseConditionerConfig
+
+__all__ = ["BaseConditioner", "BaseConditionerConfig"]

libcom/image_harmonization/source/src/lbm/models/embedders/base/base_conditioner.py

@@ -0,0 +1,60 @@
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+
+from ...base.base_model import BaseModel
+from .base_conditioner_config import BaseConditionerConfig
+
+DIM2CONDITIONING = {
+    2: "vector",
+    3: "crossattn",
+    4: "concat",
+}
+
+
+class BaseConditioner(BaseModel):
+    """This is the base class for all the conditioners. This absctacts the conditioning process
+
+    Args:
+
+        config (BaseConditionerConfig): The configuration of the conditioner
+
+    Examples
+    ########
+
+    To use the conditioner, you can import the class and use it as follows:
+
+    .. code-block:: python
+
+        from cr.models.embedders import BaseConditioner, BaseConditionerConfig
+
+        # Create the conditioner config
+        config = BaseConditionerConfig(
+            input_key="text", # The key for the input
+            unconditional_conditioning_rate=0.3, # Drops the conditioning with 30% probability during training
+        )
+
+        # Create the conditioner
+        conditioner = BaseConditioner(config)
+    """
+
+    def __init__(self, config: BaseConditionerConfig):
+        BaseModel.__init__(self, config)
+        self.config = config
+        self.input_key = config.input_key
+        self.dim2outputkey = DIM2CONDITIONING
+        self.ucg_rate = config.unconditional_conditioning_rate
+
+    def forward(
+        self, batch: Dict[str, Any], force_zero_embedding: bool = False, *args, **kwargs
+    ):
+        """
+         Forward pass of the embedder.
+
+        Args:
+
+            batch (Dict[str, Any]): A dictionary containing the input data.
+            force_zero_embedding (bool): Whether to force zero embedding.
+                This will return an embedding with all entries set to 0. Defaults to False.
+        """
+        raise NotImplementedError("Forward pass must be implemented in child class")

libcom/image_harmonization/source/src/lbm/models/embedders/base/base_conditioner_config.py

@@ -0,0 +1,27 @@
+from typing import Literal
+
+from pydantic.dataclasses import dataclass
+
+from ....config import BaseConfig
+
+
+@dataclass
+class BaseConditionerConfig(BaseConfig):
+    """This is the ClipEmbedderConfig class which defines all the useful parameters to instantiate the model
+
+    Args:
+
+        input_key (str): The key for the input. Defaults to "text".
+        unconditional_conditioning_rate (float): Drops the conditioning with this probability during training. Defaults to 0.0.
+    """
+
+    input_key: str = "text"
+    unconditional_conditioning_rate: float = 0.0
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        assert (
+            self.unconditional_conditioning_rate >= 0.0
+            and self.unconditional_conditioning_rate <= 1.0
+        ), "Unconditional conditioning rate should be between 0 and 1"

libcom/image_harmonization/source/src/lbm/models/embedders/conditioners_wrapper.py

@@ -0,0 +1,114 @@
+import logging
+from typing import Any, Dict, List, Union
+
+import torch
+import torch.nn as nn
+
+from .base import BaseConditioner
+
+KEY2CATDIM = {
+    "vector": 1,
+    "crossattn": 2,
+    "concat": 1,
+}
+
+
+class ConditionerWrapper(nn.Module):
+    """
+    Wrapper for conditioners. This class allows to apply multiple conditioners in a single forward pass.
+
+    Args:
+
+        conditioners (List[BaseConditioner]): List of conditioners to apply in the forward pass.
+    """
+
+    def __init__(
+        self,
+        conditioners: Union[List[BaseConditioner], None] = None,
+    ):
+        nn.Module.__init__(self)
+        self.conditioners = nn.ModuleList(conditioners)
+        self.device = torch.device("cpu")
+        self.dtype = torch.float32
+
+    def conditioner_sanity_check(self):
+        cond_input_keys = []
+        for conditioner in self.conditioners:
+            cond_input_keys.append(conditioner.input_key)
+
+        assert all([key in set(cond_input_keys) for key in self.ucg_keys])
+
+    def on_fit_start(self, device: torch.device | None = None, *args, **kwargs):
+        """Called when the training starts"""
+        for conditioner in self.conditioners:
+            conditioner.on_fit_start(device=device, *args, **kwargs)
+
+    def forward(
+        self,
+        batch: Dict[str, Any],
+        ucg_keys: List[str] = None,
+        set_ucg_rate_zero=False,
+        *args,
+        **kwargs,
+    ):
+        """
+        Forward pass through all conditioners
+
+        Args:
+
+            batch: batch of data
+            ucg_keys: keys to use for ucg. This will force zero conditioning in all the
+                conditioners that have input_keys in ucg_keys
+            set_ucg_rate_zero: set the ucg rate to zero for all the conditioners except the ones in ucg_keys
+
+        Returns:
+
+        Dict[str, Any]: The output of the conditioner. The output of the conditioner is a dictionary with the main key "cond" and value
+            is a dictionary with the keys as the type of conditioning and the value as the conditioning tensor.
+        """
+        if ucg_keys is None:
+            ucg_keys = []
+        wrapper_outputs = dict(cond={})
+        for conditioner in self.conditioners:
+            if conditioner.input_key in ucg_keys:
+                force_zero_embedding = True
+            elif conditioner.ucg_rate > 0 and not set_ucg_rate_zero:
+                force_zero_embedding = bool(torch.rand(1) < conditioner.ucg_rate)
+            else:
+                force_zero_embedding = False
+
+            conditioner_output = conditioner.forward(
+                batch, force_zero_embedding=force_zero_embedding, *args, **kwargs
+            )
+            logging.debug(
+                f"conditioner:{conditioner.__class__.__name__}, input_key:{conditioner.input_key}, force_ucg_zero_embedding:{force_zero_embedding}"
+            )
+            for key in conditioner_output:
+                logging.debug(
+                    f"conditioner_output:{key}:{conditioner_output[key].shape}"
+                )
+                if key in wrapper_outputs["cond"]:
+                    wrapper_outputs["cond"][key] = torch.cat(
+                        [wrapper_outputs["cond"][key], conditioner_output[key]],
+                        KEY2CATDIM[key],
+                    )
+                else:
+                    wrapper_outputs["cond"][key] = conditioner_output[key]
+
+        return wrapper_outputs
+
+    def to(self, *args, **kwargs):
+        """
+        Move all conditioners to device and dtype
+        """
+        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
+        self = super().to(device=device, dtype=dtype, non_blocking=non_blocking)
+        for conditioner in self.conditioners:
+            conditioner.to(device=device, dtype=dtype, non_blocking=non_blocking)
+
+        if device is not None:
+            self.device = device
+        if dtype is not None:
+            self.dtype = dtype
+
+        return self

libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/__init__.py

@@ -0,0 +1,4 @@
+from .latents_concat_embedder_config import LatentsConcatEmbedderConfig
+from .latents_concat_embedder_model import LatentsConcatEmbedder
+
+__all__ = ["LatentsConcatEmbedder", "LatentsConcatEmbedderConfig"]

libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/latents_concat_embedder_config.py

@@ -0,0 +1,31 @@
+from dataclasses import field
+from typing import List, Union
+
+from pydantic.dataclasses import dataclass
+
+from ..base import BaseConditionerConfig
+
+
+@dataclass
+class LatentsConcatEmbedderConfig(BaseConditionerConfig):
+    """
+    Configs for the LatentsConcatEmbedder embedder
+
+    Args:
+        image_keys (Union[List[str], None]): Keys of the images to compute the VAE embeddings
+        mask_keys (Union[List[str], None]): Keys of the masks to resize
+    """
+
+    image_keys: Union[List[str], None] = field(default_factory=lambda: ["image"])
+    mask_keys: Union[List[str], None] = field(default_factory=lambda: ["mask"])
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        # Make sure that at least one of the image_keys or mask_keys is provided
+        assert (self.image_keys is not None) or (
+            self.mask_keys is not None
+        ), "At least one of the image_keys or mask_keys must be provided."
+
+        self.image_keys = self.image_keys if self.image_keys is not None else []
+        self.mask_keys = self.mask_keys if self.mask_keys is not None else []

libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/latents_concat_embedder_model.py

@@ -0,0 +1,80 @@
+from typing import Any, Dict
+
+import torch
+import torchvision.transforms.functional as F
+
+from lbm.models.vae import AutoencoderKLDiffusers
+
+from ..base import BaseConditioner
+from .latents_concat_embedder_config import LatentsConcatEmbedderConfig
+
+
+class LatentsConcatEmbedder(BaseConditioner):
+    """
+    Class computing VAE embeddings from given images and resizing the masks.
+    Then outputs are then concatenated to the noise in the latent space.
+
+    Args:
+        config (LatentsConcatEmbedderConfig): Configs to create the embedder
+    """
+
+    def __init__(self, config: LatentsConcatEmbedderConfig):
+        BaseConditioner.__init__(self, config)
+
+    def forward(
+        self, batch: Dict[str, Any], vae: AutoencoderKLDiffusers, *args, **kwargs
+    ) -> dict:
+        """
+        Args:
+            batch (dict): A batch of images to be processed by this embedder. In the batch,
+            the images must range between [-1, 1] and the masks range between [0, 1].
+            vae (AutoencoderKLDiffusers): VAE
+
+        Returns:
+            output (dict): outputs
+        """
+
+        # Check if image are of the same size
+        dims_list = []
+        for image_key in self.config.image_keys:
+            dims_list.append(batch[image_key].shape[-2:])
+        for mask_key in self.config.mask_keys:
+            dims_list.append(batch[mask_key].shape[-2:])
+        assert all(
+            dims == dims_list[0] for dims in dims_list
+        ), "All images and masks must have the same dimensions."
+
+        # Find the latent dimensions
+        if len(self.config.image_keys) > 0:
+            latent_dims = (
+                batch[self.config.image_keys[0]].shape[-2] // vae.downsampling_factor,
+                batch[self.config.image_keys[0]].shape[-1] // vae.downsampling_factor,
+            )
+        else:
+            latent_dims = (
+                batch[self.config.mask_keys[0]].shape[-2] // vae.downsampling_factor,
+                batch[self.config.mask_keys[0]].shape[-1] // vae.downsampling_factor,
+            )
+
+        outputs = []
+
+        # Resize the masks and concat them
+        for mask_key in self.config.mask_keys:
+            curr_latents = F.resize(
+                batch[mask_key],
+                size=latent_dims,
+                interpolation=F.InterpolationMode.BILINEAR,
+            )
+            outputs.append(curr_latents)
+
+        # Compute VAE embeddings from the images
+        for image_key in self.config.image_keys:
+            vae_embs = vae.encode(batch[image_key])
+            outputs.append(vae_embs)
+
+        # Concat all the outputs
+        outputs = torch.concat(outputs, dim=1)
+
+        outputs = {self.dim2outputkey[outputs.dim()]: outputs}
+
+        return outputs

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

@@ -0,0 +1,4 @@
+from .lbm_config import LBMConfig
+from .lbm_model import LBMModel
+
+__all__ = ["LBMModel", "LBMConfig"]