diff --git a/.pylintrc b/.pylintrc index 68b004820..8ae641875 100644 --- a/.pylintrc +++ b/.pylintrc @@ -46,6 +46,8 @@ ignore-paths=/usr/lib/.*$, scripts/freescale, scripts/infiniteyou, scripts/instantir, + scripts/lbm, + scripts/layerdiffuse, scripts/mod, scripts/pixelsmith, scripts/differential_diffusion.py, diff --git a/.ruff.toml b/.ruff.toml index b21563996..1b377e615 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -22,6 +22,7 @@ exclude = [ "pipelines/omnigen2", "pipelines/segmoe", + "scripts/lbm", "scripts/xadapter", "scripts/pulid", "scripts/instantir", diff --git a/CHANGELOG.md b/CHANGELOG.md index b04e59424..081156d9c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,13 +2,19 @@ ## Update for 2025-07-04 -- **UI** - - major update to modernui layout - - redesign of the Flat UI theme - **Models** + - [LBM: Latent Bridge Matching](https://github.com/gojasper/LBM) + very fast automatic image background replacement methods with relightning! + *simple*: automatic background replacement using [BiRefNet](https://github.com/ZhengPeng7/BiRefNet) + *relighting*: automatic background replacement with reglighting so source image fits desired background + with optional composite blending + available in *img2img or control -> scripts* - Add **FLUX.1-Kontext-Dev** inpaint workflow - Support **FLUX.1** all-in-one safetensors - - Support TAESD preview and remote VAE for **HunyuanDit** + - Support TAESD preview and remote VAE for **HunyuanDit** +- **UI** + - major update to modernui layout + - redesign of the Flat UI theme - **Compute** - support for [SageAttention2++](https://github.com/thu-ml/SageAttention) provides 10-15% performance improvement over default SDPA for transformer-based models! diff --git a/TODO.md b/TODO.md index 59469e97f..bd235c148 100644 --- a/TODO.md +++ b/TODO.md @@ -26,18 +26,17 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma - [IPAdapter negative guidance](https://github.com/huggingface/diffusers/discussions/7167) - [IPAdapter composition](https://huggingface.co/ostris/ip-composition-adapter) -- [Refactor attention](https://github.com/huggingface/diffusers/pull/11311) - [STG](https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance) -- [LBM](https://github.com/gojasper/LBM) - [SmoothCache](https://github.com/huggingface/diffusers/issues/11135) - [MagCache](https://github.com/lllyasviel/FramePack/pull/673/files) - [HiDream GGUF](https://github.com/huggingface/diffusers/pull/11550) -- [Diffusers guiders](https://github.com/huggingface/diffusers/pull/11311) - [Nunchaku PulID](https://github.com/mit-han-lab/nunchaku/pull/274) - [Dream0 guidance](https://huggingface.co/ByteDance/DreamO) -- [S3Diff diffusion upscaler](https://github.com/ArcticHare105/S3Diff) - [SUPIR upscaler](https://github.com/Fanghua-Yu/SUPIR) -- [TensorRT](https://github.com/huggingface/diffusers/pull/11173) + +### Future Considerations +- [TensorRT](https://github.com/huggingface/diffusers/pull/11173) +- [Modular guiders](https://github.com/huggingface/diffusers/pull/11311) ### New models diff --git a/modules/model_quant.py b/modules/model_quant.py index 6bf36ed0e..2933fb659 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -314,7 +314,7 @@ def load_fp8_model_layerwise(checkpoint_info, load_model_func, diffusers_load_co load_args["torch_dtype"] = storage_dtype model = load_model_func(repo_path, **load_args) model = upcast_non_layerwise_modules(model, devices.dtype) - model._skip_layerwise_casting_patterns = None + model._skip_layerwise_casting_patterns = None # pylint: disable=protected-access model.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=False, skip_modules_pattern=[]) model.layerwise_storage_dtype = storage_dtype model.quantization_method = 'LayerWise' @@ -345,7 +345,7 @@ def apply_layerwise(sd_model, quiet:bool=False): try: cls = getattr(sd_model, module).__class__.__name__ m = getattr(sd_model, module) - if getattr(m, "quantization_method", None) in {'LayerWise', quantization_config.QuantizationMethod.LAYERWISE}: + if getattr(m, "quantization_method", None) in {'LayerWise', quantization_config.QuantizationMethod.LAYERWISE}: # pylint: disable=no-member storage_dtype = getattr(m, "layerwise_storage_dtype", storage_dtype) m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking) elif module.startswith('unet') and ('Model' in shared.opts.layerwise_quantization): diff --git a/modules/rife/loss.py b/modules/rife/loss.py index 993f319fb..8b6309006 100644 --- a/modules/rife/loss.py +++ b/modules/rife/loss.py @@ -101,8 +101,7 @@ class VGGPerceptualLoss(torch.nn.Module): pretrained = True self.vgg_pretrained_features = models.vgg19( pretrained=pretrained).features - self.normalize = MeanShift([0.485, 0.456, 0.406], [ - 0.229, 0.224, 0.225], norm=True).cuda() + self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).to(device=devices.device) for param in self.parameters(): param.requires_grad = False diff --git a/modules/scripts.py b/modules/scripts.py index 1de7b0a1e..49debeac8 100644 --- a/modules/scripts.py +++ b/modules/scripts.py @@ -1,2 +1,2 @@ # compatibility with extensions that import scripts directly -from modules.scripts_manager import * # noqa: F403 +from modules.scripts_manager import * # noqa: F403 # pylint: disable=wildcard-import diff --git a/modules/sd_models.py b/modules/sd_models.py index df1758928..f02418d35 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1168,7 +1168,7 @@ def hf_auth_check(checkpoint_info): try: login = modelloader.hf_login() repo_id = path_to_repo(checkpoint_info) - hf.auth_check(repo_id) + return hf.auth_check(repo_id) except Exception as e: shared.log.error(f'Load model: repo="{repo_id}" login={login} {e}') return False diff --git a/pipelines/model_chroma.py b/pipelines/model_chroma.py index 2c03af8c5..a8c5bad42 100644 --- a/pipelines/model_chroma.py +++ b/pipelines/model_chroma.py @@ -4,8 +4,8 @@ import torch import diffusers import transformers from safetensors.torch import load_file -from huggingface_hub import hf_hub_download, auth_check -from modules import shared, errors, devices, modelloader, sd_models, sd_unet, model_te, model_quant, sd_hijack_te +from huggingface_hub import hf_hub_download +from modules import shared, errors, devices, sd_models, sd_unet, model_te, model_quant, sd_hijack_te debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -110,7 +110,7 @@ def load_chroma_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=u return transformer, text_encoder -def load_quants(kwargs, repo_id, cache_dir, allow_quant): +def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unused-argument try: diffusers_load_config = { "torch_dtype": devices.dtype, diff --git a/pipelines/model_cosmos.py b/pipelines/model_cosmos.py index 87b85db0d..c9fad08db 100644 --- a/pipelines/model_cosmos.py +++ b/pipelines/model_cosmos.py @@ -1,7 +1,7 @@ import os import transformers import diffusers -from modules import shared, devices, sd_models, model_quant, modelloader, sd_hijack_te +from modules import shared, devices, sd_models, model_quant, sd_hijack_te def load_transformer(repo_id, diffusers_load_config={}): diff --git a/pipelines/model_flex.py b/pipelines/model_flex.py index f531dac79..4a11152ff 100644 --- a/pipelines/model_flex.py +++ b/pipelines/model_flex.py @@ -1,7 +1,7 @@ import os import transformers import diffusers -from modules import shared, devices, sd_models, model_quant, modelloader, sd_hijack_te +from modules import shared, devices, sd_models, model_quant, sd_hijack_te def load_transformer(repo_id, diffusers_load_config={}): diff --git a/pipelines/model_flux.py b/pipelines/model_flux.py index 38dafa3ff..77c0e768e 100644 --- a/pipelines/model_flux.py +++ b/pipelines/model_flux.py @@ -5,7 +5,7 @@ import diffusers import transformers from safetensors.torch import load_file from huggingface_hub import hf_hub_download -from modules import shared, errors, devices, modelloader, sd_models, sd_unet, model_te, model_quant, sd_hijack_te +from modules import shared, errors, devices, sd_models, sd_unet, model_te, model_quant, sd_hijack_te debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -108,7 +108,7 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu return transformer, text_encoder_2 -def load_quants(kwargs, repo_id, cache_dir, allow_quant): +def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unused-argument try: diffusers_load_config = { "torch_dtype": devices.dtype, diff --git a/pipelines/model_hidream.py b/pipelines/model_hidream.py index 50ff59fc0..6df9092e2 100644 --- a/pipelines/model_hidream.py +++ b/pipelines/model_hidream.py @@ -1,7 +1,7 @@ import os import transformers import diffusers -from modules import shared, devices, sd_models, model_quant, modelloader, sd_hijack_te +from modules import shared, devices, sd_models, model_quant, sd_hijack_te def load_transformer(repo_id, diffusers_load_config={}): diff --git a/pipelines/model_sd3.py b/pipelines/model_sd3.py index f43b74a54..8b3b6601c 100644 --- a/pipelines/model_sd3.py +++ b/pipelines/model_sd3.py @@ -1,7 +1,7 @@ import os import diffusers import transformers -from modules import shared, devices, errors, sd_models, sd_unet, model_quant, model_tools, modelloader +from modules import shared, devices, errors, sd_models, sd_unet, model_quant, model_tools def load_overrides(kwargs, cache_dir): diff --git a/scripts/freescale/free_lunch_utils.py b/scripts/freescale/free_lunch_utils.py index be26b732a..ebf165105 100644 --- a/scripts/freescale/free_lunch_utils.py +++ b/scripts/freescale/free_lunch_utils.py @@ -2,6 +2,7 @@ from typing import Any, Dict, Optional, Tuple import torch import torch.fft as fft from diffusers.utils import is_torch_version +from modules import devices """ Borrowed from https://github.com/ChenyangSi/FreeU/blob/main/demo/free_lunch_utils.py """ @@ -29,7 +30,7 @@ def Fourier_filter(x, threshold, scale): x_freq = fft.fftshift(x_freq, dim=(-2, -1)) B, C, H, W = x_freq.shape - mask = torch.ones((B, C, H, W)).cuda() + mask = torch.ones((B, C, H, W)).to(device=devices.device) crow, ccol = H // 2, W //2 mask[..., crow - threshold:crow + threshold, ccol - threshold:ccol + threshold] = scale diff --git a/scripts/infiniteyou/pipeline_infu_flux.py b/scripts/infiniteyou/pipeline_infu_flux.py index fc8b84647..7499ab077 100644 --- a/scripts/infiniteyou/pipeline_infu_flux.py +++ b/scripts/infiniteyou/pipeline_infu_flux.py @@ -101,7 +101,7 @@ def extract_arcface_bgr_embedding(in_image, landmark, arcface_model=None, in_set arc_face_image = face_align.norm_crop(in_image, landmark=np.array(kps), image_size=112) arc_face_image = torch.from_numpy(arc_face_image).unsqueeze(0).permute(0,3,1,2) / 255. arc_face_image = 2 * arc_face_image - 1 - arc_face_image = arc_face_image.cuda().contiguous() + arc_face_image = arc_face_image.to(device=devices.device).contiguous() if arcface_model is None: arcface_model = init_recognition_model('arcface', device=devices.device) face_emb = arcface_model(arc_face_image)[0] # [512], normalized @@ -252,7 +252,7 @@ class InfUFluxPipeline: face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face landmark = face_info['kps'] id_embed = extract_arcface_bgr_embedding(id_image_cv2, landmark, self.arcface_model) - id_embed = id_embed.clone().unsqueeze(0).float().cuda() + id_embed = id_embed.clone().unsqueeze(0).float() id_embed = id_embed.reshape([1, -1, 512]) id_embed = id_embed.to(device=devices.device, dtype=devices.dtype) with torch.no_grad(): diff --git a/scripts/layerdiffuse_ext.py b/scripts/layerdiffuse_ext.py index 611af5a77..4a25e8c1b 100644 --- a/scripts/layerdiffuse_ext.py +++ b/scripts/layerdiffuse_ext.py @@ -11,7 +11,7 @@ class Script(scripts_manager.Script): return True if shared.native else False def apply(self): - from scripts import layerdiffuse + from scripts import layerdiffuse # pylint: disable=no-name-in-module if not shared.sd_loaded: shared.log.error('LayerDiffuse: model not loaded') return self.is_active() diff --git a/scripts/lbm/__init__.py b/scripts/lbm/__init__.py new file mode 100644 index 000000000..b8cbc2de6 --- /dev/null +++ b/scripts/lbm/__init__.py @@ -0,0 +1,6 @@ +from .inference import evaluate +from .utils import get_model +from .extract import extract_object, resize_and_center_crop + + +__all__ = ["evaluate", "get_model", "extract_object", "resize_and_center_crop"] diff --git a/scripts/lbm/base/__init__.py b/scripts/lbm/base/__init__.py new file mode 100644 index 000000000..205b4f7a2 --- /dev/null +++ b/scripts/lbm/base/__init__.py @@ -0,0 +1,5 @@ +from .base_model import BaseModel +from .model_config import ModelConfig + + +__all__ = ["BaseModel", "ModelConfig"] diff --git a/scripts/lbm/base/base_model.py b/scripts/lbm/base/base_model.py new file mode 100644 index 000000000..5f6778fe3 --- /dev/null +++ b/scripts/lbm/base/base_model.py @@ -0,0 +1,64 @@ +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 diff --git a/scripts/lbm/base/model_config.py b/scripts/lbm/base/model_config.py new file mode 100644 index 000000000..427b69976 --- /dev/null +++ b/scripts/lbm/base/model_config.py @@ -0,0 +1,7 @@ +from pydantic.dataclasses import dataclass +from ..config import BaseConfig + + +@dataclass +class ModelConfig(BaseConfig): + input_key: str = "image" diff --git a/scripts/lbm/config.py b/scripts/lbm/config.py new file mode 100644 index 000000000..de9948158 --- /dev/null +++ b/scripts/lbm/config.py @@ -0,0 +1,140 @@ +import json +import os +import warnings +from dataclasses import asdict, field +from typing import Any, Dict, Union +import yaml +from pydantic import ValidationError +from pydantic.dataclasses import dataclass +from yaml import safe_load + + +@dataclass +class BaseConfig: + """This is the BaseConfig class which defines all the useful loading and saving methods + of the configs""" + + name: str = field(init=False) + + def __post_init__(self): + self.name = self.__class__.__name__ + + @classmethod + def from_dict(cls, config_dict: Dict[str, Any]) -> "BaseConfig": + """Creates a BaseConfig instance from a dictionnary + + Args: + config_dict (dict): The Python dictionnary containing all the parameters + + Returns: + :class:`BaseConfig`: The created instance + """ + try: + config = cls(**config_dict) + except (ValidationError, TypeError) as e: + raise e + return config + + @classmethod + def _dict_from_json(cls, json_path: Union[str, os.PathLike]) -> Dict[str, Any]: + try: + with open(json_path) as f: + try: + config_dict = json.load(f) + return config_dict + + except (TypeError, json.JSONDecodeError) as e: + raise TypeError( + f"File {json_path} not loadable. Maybe not json ? \n" + f"Catch Exception {type(e)} with message: " + str(e) + ) from e + + except FileNotFoundError: + raise FileNotFoundError( + f"Config file not found. Please check path '{json_path}'" + ) + + @classmethod + def from_json(cls, json_path: str) -> "BaseConfig": + """Creates a BaseConfig instance from a JSON config file + + Args: + json_path (str): The path to the json file containing all the parameters + + Returns: + :class:`BaseConfig`: The created instance + """ + config_dict = cls._dict_from_json(json_path) + + config_name = config_dict.pop("name") + + if cls.__name__ != config_name: + warnings.warn( + f"You are trying to load a " + f"`{ cls.__name__}` while a " + f"`{config_name}` is given." + ) + + return cls.from_dict(config_dict) + + def to_dict(self) -> dict: + """Transforms object into a Python dictionnary + + Returns: + (dict): The dictionnary containing all the parameters""" + return asdict(self) + + def to_json_string(self): + """Transforms object into a JSON string + + Returns: + (str): The JSON str containing all the parameters""" + return json.dumps(self.to_dict()) + + def save_json(self, file_path: str): + """Saves a ``.json`` file from the dataclass + + Args: + file_path (str): path to the file + """ + with open(os.path.join(file_path), "w", encoding="utf-8") as fp: + fp.write(self.to_json_string()) + + def save_yaml(self, file_path: str): + """Saves a ``.yaml`` file from the dataclass + + Args: + file_path (str): path to the file + """ + with open(os.path.join(file_path), "w", encoding="utf-8") as fp: + yaml.dump(self.to_dict(), fp) + + @classmethod + def from_yaml(cls, yaml_path: str) -> "BaseConfig": + """Creates a BaseConfig instance from a YAML config file + + Args: + yaml_path (str): The path to the yaml file containing all the parameters + + Returns: + :class:`BaseConfig`: The created instance + """ + with open(yaml_path, "r") as f: + try: + config_dict = safe_load(f) + except yaml.YAMLError as e: + raise yaml.YAMLError( + f"File {yaml_path} not loadable. Maybe not yaml ? \n" + f"Catch Exception {type(e)} with message: " + str(e) + ) from e + + config_name = config_dict.pop("name") + + if cls.__name__ != config_name: + warnings.warn( + f"You are trying to load a " + f"`{ cls.__name__}` while a " + f"`{config_name}` is given." + ) + + return cls.from_dict(config_dict) diff --git a/scripts/lbm/embedders/__init__.py b/scripts/lbm/embedders/__init__.py new file mode 100644 index 000000000..100c51019 --- /dev/null +++ b/scripts/lbm/embedders/__init__.py @@ -0,0 +1,5 @@ +from .conditioners_wrapper import ConditionerWrapper +from .latents_concat import LatentsConcatEmbedder, LatentsConcatEmbedderConfig + + +__all__ = ["LatentsConcatEmbedder", "LatentsConcatEmbedderConfig", "ConditionerWrapper"] diff --git a/scripts/lbm/embedders/base/__init__.py b/scripts/lbm/embedders/base/__init__.py new file mode 100644 index 000000000..b9f93f6be --- /dev/null +++ b/scripts/lbm/embedders/base/__init__.py @@ -0,0 +1,5 @@ +from .base_conditioner import BaseConditioner +from .base_conditioner_config import BaseConditionerConfig + + +__all__ = ["BaseConditioner", "BaseConditionerConfig"] diff --git a/scripts/lbm/embedders/base/base_conditioner.py b/scripts/lbm/embedders/base/base_conditioner.py new file mode 100644 index 000000000..4f6893902 --- /dev/null +++ b/scripts/lbm/embedders/base/base_conditioner.py @@ -0,0 +1,58 @@ +from typing import Any, Dict +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") diff --git a/scripts/lbm/embedders/base/base_conditioner_config.py b/scripts/lbm/embedders/base/base_conditioner_config.py new file mode 100644 index 000000000..5a2eab0ee --- /dev/null +++ b/scripts/lbm/embedders/base/base_conditioner_config.py @@ -0,0 +1,24 @@ +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" diff --git a/scripts/lbm/embedders/conditioners_wrapper.py b/scripts/lbm/embedders/conditioners_wrapper.py new file mode 100644 index 000000000..de184153a --- /dev/null +++ b/scripts/lbm/embedders/conditioners_wrapper.py @@ -0,0 +1,112 @@ +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, *args, **kwargs): + 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 diff --git a/scripts/lbm/embedders/latents_concat/__init__.py b/scripts/lbm/embedders/latents_concat/__init__.py new file mode 100644 index 000000000..0151758d3 --- /dev/null +++ b/scripts/lbm/embedders/latents_concat/__init__.py @@ -0,0 +1,5 @@ +from .latents_concat_embedder_config import LatentsConcatEmbedderConfig +from .latents_concat_embedder_model import LatentsConcatEmbedder + + +__all__ = ["LatentsConcatEmbedder", "LatentsConcatEmbedderConfig"] diff --git a/scripts/lbm/embedders/latents_concat/latents_concat_embedder_config.py b/scripts/lbm/embedders/latents_concat/latents_concat_embedder_config.py new file mode 100644 index 000000000..0a60678c9 --- /dev/null +++ b/scripts/lbm/embedders/latents_concat/latents_concat_embedder_config.py @@ -0,0 +1,29 @@ +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 [] diff --git a/scripts/lbm/embedders/latents_concat/latents_concat_embedder_model.py b/scripts/lbm/embedders/latents_concat/latents_concat_embedder_model.py new file mode 100644 index 000000000..447a6f8c0 --- /dev/null +++ b/scripts/lbm/embedders/latents_concat/latents_concat_embedder_model.py @@ -0,0 +1,77 @@ +from typing import Any, Dict +import torch +import torchvision.transforms.functional as F +from ...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 diff --git a/scripts/lbm/extract.py b/scripts/lbm/extract.py new file mode 100644 index 000000000..1f91272da --- /dev/null +++ b/scripts/lbm/extract.py @@ -0,0 +1,42 @@ +import torch +from PIL import Image +from torchvision import transforms +from modules import devices + + +def extract_object(birefnet, img): + # Data settings + image_size = (1024, 1024) + transform_image = transforms.Compose( + [ + transforms.Resize(image_size), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ] + ) + + image = img + input_images = transform_image(image).unsqueeze(0).to(dtype=torch.float32, device=devices.device) + + # Prediction + with torch.no_grad(): + preds = birefnet(input_images)[-1].sigmoid().cpu() + pred = preds[0].squeeze() + pred_pil = transforms.ToPILImage()(pred) + mask = pred_pil.resize(image.size) + image = Image.composite(image, Image.new("RGB", image.size, (127, 127, 127)), mask) + return image, mask + + +def resize_and_center_crop(image, target_width, target_height): + original_width, original_height = image.size + scale_factor = max(target_width / original_width, target_height / original_height) + resized_width = int(round(original_width * scale_factor)) + resized_height = int(round(original_height * scale_factor)) + resized_image = image.resize((resized_width, resized_height), Image.Resampling.LANCZOS) + left = (resized_width - target_width) / 2 + top = (resized_height - target_height) / 2 + right = (resized_width + target_width) / 2 + bottom = (resized_height + target_height) / 2 + cropped_image = resized_image.crop((left, top, right, bottom)) + return cropped_image diff --git a/scripts/lbm/inference.py b/scripts/lbm/inference.py new file mode 100644 index 000000000..f8ec2126b --- /dev/null +++ b/scripts/lbm/inference.py @@ -0,0 +1,70 @@ +import logging +import PIL +import torch +from torchvision.transforms import ToPILImage, ToTensor +from .lbm import LBMModel +from modules import devices + + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +ASPECT_RATIOS = { + str(512 / 2048): (512, 2048), + str(1024 / 1024): (1024, 1024), + str(2048 / 512): (2048, 512), + str(896 / 1152): (896, 1152), + str(1152 / 896): (1152, 896), + str(512 / 1920): (512, 1920), + str(640 / 1536): (640, 1536), + str(768 / 1280): (768, 1280), + str(1280 / 768): (1280, 768), + str(1536 / 640): (1536, 640), + str(1920 / 512): (1920, 512), +} + + +@torch.no_grad() +def evaluate( + model: LBMModel, + source_image: PIL.Image.Image, + num_sampling_steps: int = 1, +): + """ + Evaluate the model on an image coming from the source distribution and generate a new image from the target distribution. + + Args: + model (LBMModel): The model to evaluate. + source_image (PIL.Image.Image): The source image to evaluate the model on. + num_sampling_steps (int): The number of sampling steps to use for the model. + + Returns: + PIL.Image.Image: The generated image. + """ + + ori_h_bg, ori_w_bg = source_image.size + ar_bg = ori_h_bg / ori_w_bg + closest_ar_bg = min(ASPECT_RATIOS, key=lambda x: abs(float(x) - ar_bg)) + source_dimensions = ASPECT_RATIOS[closest_ar_bg] + + source_image = source_image.resize(source_dimensions) + + img_pasted_tensor = ToTensor()(source_image).unsqueeze(0) * 2 - 1 + batch = { + "source_image": img_pasted_tensor.to(dtype=devices.dtype, device=devices.device), + } + + z_source = model.vae.encode(batch[model.source_key]) + + output_image = model.sample( + z=z_source, + num_steps=num_sampling_steps, + conditioner_inputs=batch, + max_samples=1, + ).clamp(-1, 1) + + output_image = (output_image[0].float().cpu() + 1) / 2 + output_image = ToPILImage()(output_image) + output_image.resize((ori_h_bg, ori_w_bg)) + + return output_image diff --git a/scripts/lbm/lbm/__init__.py b/scripts/lbm/lbm/__init__.py new file mode 100644 index 000000000..e907b0c91 --- /dev/null +++ b/scripts/lbm/lbm/__init__.py @@ -0,0 +1,5 @@ +from .lbm_config import LBMConfig +from .lbm_model import LBMModel + + +__all__ = ["LBMModel", "LBMConfig"] diff --git a/scripts/lbm/lbm/lbm_config.py b/scripts/lbm/lbm/lbm_config.py new file mode 100644 index 000000000..a9b1b40ac --- /dev/null +++ b/scripts/lbm/lbm/lbm_config.py @@ -0,0 +1,99 @@ +from typing import List, Literal, Optional +from pydantic.dataclasses import dataclass +from ..base import ModelConfig + + +@dataclass +class LBMConfig(ModelConfig): + """This is the Config for LBM Model class which defines all the useful parameters to be used in the model. + + Args: + + source_key (str): + Key for the source image. Defaults to "source_image" + + target_key (str): + Key for the target image. Defaults to "target_image" + + mask_key (Optional[str]): + Key for the mask showing the valid pixels. Defaults to None + + latent_loss_type (str): + Loss type to use. Defaults to "l2". Choices are "l2", "l1" + + pixel_loss_type (str): + Pixel loss type to use. Defaults to "l2". Choices are "l2", "l1", "lpips" + + pixel_loss_max_size (int): + Maximum size of the image for pixel loss. + The image will be cropped to this size to reduce decoding computation cost. Defaults to 512 + + pixel_loss_weight (float): + Weight of the pixel loss. Defaults to 0.0 + + timestep_sampling (str): + Timestep sampling to use. Defaults to "uniform". Choices are "uniform" + + input_key (str): + Key for the input. Defaults to "image" + + controlnet_input_key (str): + Key for the controlnet conditioning. Defaults to "controlnet_conditioning" + + adapter_input_key (str): + Key for the adapter conditioning. Defaults to "adapter_conditioning" + + ucg_keys (Optional[List[str]]): + List of keys for which we enforce zero_conditioning during Classifier-free guidance. Defaults to None + + prediction_type (str): + Type of prediction to use. Defaults to "epsilon". Choices are "epsilon", "v_prediction", "flow + + logit_mean (Optional[float]): + Mean of the logit for the log normal distribution. Defaults to 0.0 + + logit_std (Optional[float]): + Standard deviation of the logit for the log normal distribution. Defaults to 1.0 + + guidance_scale (Optional[float]): + The guidance scale. Useful for finetunning guidance distilled diffusion models. Defaults to None + + selected_timesteps (Optional[List[float]]): + List of selected timesteps to be sampled from if using `custom_timesteps` timestep sampling. Defaults to None + + prob (Optional[List[float]]): + List of probabilities for the selected timesteps if using `custom_timesteps` timestep sampling. Defaults to None + """ + + source_key: str = "source_image" + target_key: str = "target_image" + mask_key: Optional[str] = None + latent_loss_weight: float = 1.0 + latent_loss_type: Literal["l2", "l1"] = "l2" + pixel_loss_type: Literal["l2", "l1", "lpips"] = "l2" + pixel_loss_max_size: int = 512 + pixel_loss_weight: float = 0.0 + timestep_sampling: Literal["uniform", "log_normal", "custom_timesteps"] = "uniform" + logit_mean: Optional[float] = 0.0 + logit_std: Optional[float] = 1.0 + selected_timesteps: Optional[List[float]] = None + prob: Optional[List[float]] = None + bridge_noise_sigma: float = 0.001 + + def __post_init__(self): + super().__post_init__() + if self.timestep_sampling == "log_normal": + assert isinstance(self.logit_mean, float) and isinstance( + self.logit_std, float + ), "logit_mean and logit_std should be float for log_normal timestep sampling" + + if self.timestep_sampling == "custom_timesteps": + assert isinstance(self.selected_timesteps, list) and isinstance( + self.prob, list + ), "timesteps and prob should be list for custom_timesteps timestep sampling" + assert len(self.selected_timesteps) == len( + self.prob + ), "timesteps and prob should be of same length for custom_timesteps timestep sampling" + assert ( + sum(self.prob) == 1 + ), "prob should sum to 1 for custom_timesteps timestep sampling" diff --git a/scripts/lbm/lbm/lbm_model.py b/scripts/lbm/lbm/lbm_model.py new file mode 100644 index 000000000..ac4f63332 --- /dev/null +++ b/scripts/lbm/lbm/lbm_model.py @@ -0,0 +1,509 @@ +from typing import Any, Dict, List, Optional, Tuple, Union +import lpips +import numpy as np +import torch +import torch.nn as nn +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from tqdm import tqdm +from ..base.base_model import BaseModel +from ..embedders import ConditionerWrapper +from ..unets import DiffusersUNet2DCondWrapper, DiffusersUNet2DWrapper +from ..vae import AutoencoderKLDiffusers +from .lbm_config import LBMConfig + + +class LBMModel(BaseModel): + """This is the LBM class which defines the model. + + Args: + + config (LBMConfig): + Configuration for the model + + denoiser (Union[DiffusersUNet2DWrapper, DiffusersTransformer2DWrapper]): + Denoiser to use for the diffusion model. Defaults to None + + training_noise_scheduler (EulerDiscreteScheduler): + Noise scheduler to use for training. Defaults to None + + sampling_noise_scheduler (EulerDiscreteScheduler): + Noise scheduler to use for sampling. Defaults to None + + vae (AutoencoderKLDiffusers): + VAE to use for the diffusion model. Defaults to None + + conditioner (ConditionerWrapper): + Conditioner to use for the diffusion model. Defaults to None + """ + + @classmethod + def load_from_config(cls, config: LBMConfig): + return cls(config=config) + + def __init__( + self, + config: LBMConfig, + denoiser: Union[ + DiffusersUNet2DWrapper, + DiffusersUNet2DCondWrapper, + ] = None, + training_noise_scheduler: FlowMatchEulerDiscreteScheduler = None, + sampling_noise_scheduler: FlowMatchEulerDiscreteScheduler = None, + vae: AutoencoderKLDiffusers = None, + conditioner: ConditionerWrapper = None, + ): + BaseModel.__init__(self, config) + + self.vae = vae + self.denoiser = denoiser + self.conditioner = conditioner + self.sampling_noise_scheduler = sampling_noise_scheduler + self.training_noise_scheduler = training_noise_scheduler + self.timestep_sampling = config.timestep_sampling + self.latent_loss_type = config.latent_loss_type + self.latent_loss_weight = config.latent_loss_weight + self.pixel_loss_type = config.pixel_loss_type + self.pixel_loss_max_size = config.pixel_loss_max_size + self.pixel_loss_weight = config.pixel_loss_weight + self.logit_mean = config.logit_mean + self.logit_std = config.logit_std + self.prob = config.prob + self.selected_timesteps = config.selected_timesteps + self.source_key = config.source_key + self.target_key = config.target_key + self.mask_key = config.mask_key + self.bridge_noise_sigma = config.bridge_noise_sigma + + self.num_iterations = nn.Parameter( + torch.tensor(0, dtype=torch.float32), requires_grad=False + ) + if self.pixel_loss_type == "lpips" and self.pixel_loss_weight > 0: + self.lpips_loss = lpips.LPIPS(net="vgg") + + else: + self.lpips_loss = None + + def on_fit_start(self, device: torch.device | None = None, *args, **kwargs): + """Called when the training starts""" + super().on_fit_start(device=device, *args, **kwargs) + if self.vae is not None: + self.vae.on_fit_start(device=device, *args, **kwargs) + if self.conditioner is not None: + self.conditioner.on_fit_start(device=device, *args, **kwargs) + + def forward(self, batch: Dict[str, Any], step=0, batch_idx=0, *args, **kwargs): + + self.num_iterations += 1 + + # Get inputs/latents + if self.vae is not None: + vae_inputs = batch[self.target_key] + z = self.vae.encode(vae_inputs) + downsampling_factor = self.vae.downsampling_factor + else: + z = batch[self.target_key] + downsampling_factor = 1 + + if self.mask_key in batch: + valid_mask = batch[self.mask_key].bool()[:, 0, :, :].unsqueeze(1) + invalid_mask = ~valid_mask + valid_mask_for_latent = ~torch.max_pool2d( + invalid_mask.float(), + downsampling_factor, + downsampling_factor, + ).bool() + valid_mask_for_latent = valid_mask_for_latent.repeat((1, z.shape[1], 1, 1)) + + else: + valid_mask = torch.ones_like(batch[self.target_key]).bool() + valid_mask_for_latent = torch.ones_like(z).bool() + + source_image = batch[self.source_key] + source_image = torch.nn.functional.interpolate( + source_image, + size=batch[self.target_key].shape[-2:], + mode="bilinear", + align_corners=False, + ).to(z.dtype) + if self.vae is not None: + z_source = self.vae.encode(source_image) + + else: + z_source = source_image + + # Get conditionings + conditioning = self._get_conditioning(batch, *args, **kwargs) + + # Sample a timestep + timestep = self._timestep_sampling(n_samples=z.shape[0], device=z.device) + sigmas = None + + # Create interpolant + sigmas = self._get_sigmas( + self.training_noise_scheduler, timestep, n_dim=4, device=z.device + ) + noisy_sample = ( + sigmas * z_source + + (1.0 - sigmas) * z + + self.bridge_noise_sigma + * (sigmas * (1.0 - sigmas)) ** 0.5 + * torch.randn_like(z) + ) + + for i, t in enumerate(timestep): + if t.item() == self.training_noise_scheduler.timesteps[0]: + noisy_sample[i] = z_source[i] + + # Predict noise level using denoiser + prediction = self.denoiser( + sample=noisy_sample, + timestep=timestep, + conditioning=conditioning, + *args, + **kwargs, + ) + + target = z_source - z + denoised_sample = noisy_sample - prediction * sigmas + target_pixels = batch[self.target_key] + + # Compute loss + if self.latent_loss_weight > 0: + loss = self.latent_loss(prediction, target.detach(), valid_mask_for_latent) + latent_recon_loss = loss.mean() + + else: + loss = torch.zeros(z.shape[0], device=z.device) + latent_recon_loss = torch.zeros_like(loss) + + if self.pixel_loss_weight > 0: + denoised_sample = self._predicted_x_0( + model_output=prediction, + sample=noisy_sample, + sigmas=sigmas, + ) + pixel_loss = self.pixel_loss( + denoised_sample, target_pixels.detach(), valid_mask + ) + loss += self.pixel_loss_weight * pixel_loss + + else: + pixel_loss = torch.zeros_like(latent_recon_loss) + + return { + "loss": loss.mean(), + "latent_recon_loss": latent_recon_loss, + "pixel_recon_loss": pixel_loss.mean(), + "predicted_hr": denoised_sample, + "noisy_sample": noisy_sample, + } + + def latent_loss(self, prediction, model_input, valid_latent_mask): + if self.latent_loss_type == "l2": + return torch.mean( + ( + (prediction * valid_latent_mask - model_input * valid_latent_mask) + ** 2 + ).reshape(model_input.shape[0], -1), + 1, + ) + elif self.latent_loss_type == "l1": + return torch.mean( + torch.abs( + prediction * valid_latent_mask - model_input * valid_latent_mask + ).reshape(model_input.shape[0], -1), + 1, + ) + else: + raise NotImplementedError( + f"Loss type {self.latent_loss_type} not implemented" + ) + + def pixel_loss(self, prediction, model_input, valid_mask): + + latent_crop = self.pixel_loss_max_size // self.vae.downsampling_factor + input_crop = self.pixel_loss_max_size + + crop_h = max((prediction.shape[2] - latent_crop), 0) + crop_w = max((prediction.shape[3] - latent_crop), 0) + + input_crop_h = max((model_input.shape[2] - self.pixel_loss_max_size), 0) + input_crop_w = max((model_input.shape[3] - self.pixel_loss_max_size), 0) + + # image random cropping + if crop_h == 0: + offset_h = 0 + else: + offset_h = torch.randint(0, crop_h, (1,)).item() + + if crop_w == 0: + offset_w = 0 + else: + offset_w = torch.randint(0, crop_w, (1,)).item() + input_offset_h = offset_h * self.vae.downsampling_factor + input_offset_w = offset_w * self.vae.downsampling_factor + + prediction = prediction[ + :, + :, + crop_h + - offset_h : min(crop_h - offset_h + latent_crop, prediction.shape[2]), + crop_w + - offset_w : min(crop_w - offset_w + latent_crop, prediction.shape[3]), + ] + + model_input = model_input[ + :, + :, + input_crop_h + - input_offset_h : min( + input_crop_h - input_offset_h + input_crop, model_input.shape[2] + ), + input_crop_w + - input_offset_w : min( + input_crop_w - input_offset_w + input_crop, model_input.shape[3] + ), + ] + + valid_mask = valid_mask[ + :, + :, + input_crop_h + - input_offset_h : min( + input_crop_h - input_offset_h + input_crop, valid_mask.shape[2] + ), + input_crop_w + - input_offset_w : min( + input_crop_w - input_offset_w + input_crop, valid_mask.shape[3] + ), + ] + + decoded_prediction = self.vae.decode(prediction).clamp(-1, 1) + + if self.pixel_loss_type == "l2": + return torch.mean( + ( + (decoded_prediction * valid_mask - model_input * valid_mask) ** 2 + ).reshape(model_input.shape[0], -1), + 1, + ) + + elif self.pixel_loss_type == "l1": + return torch.mean( + torch.abs( + decoded_prediction * valid_mask - model_input * valid_mask + ).reshape(model_input.shape[0], -1), + 1, + ) + + elif self.pixel_loss_type == "lpips": + return self.lpips_loss( + decoded_prediction * valid_mask, model_input * valid_mask + ).mean() + + def _get_conditioning( + self, + batch: Dict[str, Any], + ucg_keys: List[str] = None, + set_ucg_rate_zero=False, + *args, + **kwargs, + ): + """ + Get the conditionings + """ + if self.conditioner is not None: + return self.conditioner( + batch, + ucg_keys=ucg_keys, + set_ucg_rate_zero=set_ucg_rate_zero, + vae=self.vae, + *args, + **kwargs, + ) + else: + return None + + def _timestep_sampling(self, n_samples=1, device="cpu"): + if self.timestep_sampling == "uniform": + idx = torch.randint( + 0, + self.training_noise_scheduler.config.num_train_timesteps, + (n_samples,), + device="cpu", + ) + return self.training_noise_scheduler.timesteps[idx].to(device=device) + + elif self.timestep_sampling == "log_normal": + u = torch.normal( + mean=self.logit_mean, + std=self.logit_std, + size=(n_samples,), + device="cpu", + ) + u = torch.nn.functional.sigmoid(u) + indices = ( + u * self.training_noise_scheduler.config.num_train_timesteps + ).long() + return self.training_noise_scheduler.timesteps[indices].to(device=device) + + elif self.timestep_sampling == "custom_timesteps": + idx = np.random.choice(len(self.selected_timesteps), n_samples, p=self.prob) + + return torch.tensor( + self.selected_timesteps, device=device, dtype=torch.long + )[idx] + + def _predicted_x_0( + self, + model_output, + sample, + sigmas=None, + ): + """ + Predict x_0, the orinal denoised sample, using the model output and the timesteps depending on the prediction type. + """ + pred_x_0 = sample - model_output * sigmas + return pred_x_0 + + def _get_sigmas( + self, scheduler, timesteps, n_dim=4, dtype=torch.float32, device="cpu" + ): + sigmas = scheduler.sigmas.to(device=device, dtype=dtype) + schedule_timesteps = scheduler.timesteps.to(device) + timesteps = timesteps.to(device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + @torch.no_grad() + def sample( + self, + z: torch.Tensor, + num_steps: int = 20, + conditioner_inputs: Optional[Dict[str, Any]] = None, + max_samples: Optional[int] = None, + verbose: bool = False, + ): + self.sampling_noise_scheduler.set_timesteps( + sigmas=np.linspace(1, 1 / num_steps, num_steps) + ) + + sample = z + + # Get conditioning + conditioning = self._get_conditioning( + conditioner_inputs, set_ucg_rate_zero=True, device=z.device + ) + + # If max_samples parameter is provided, limit the number of samples + if max_samples is not None: + sample = sample[:max_samples] + + if conditioning: + conditioning["cond"] = { + k: v[:max_samples] for k, v in conditioning["cond"].items() + } + + for i, t in tqdm( + enumerate(self.sampling_noise_scheduler.timesteps), disable=not verbose + ): + if hasattr(self.sampling_noise_scheduler, "scale_model_input"): + denoiser_input = self.sampling_noise_scheduler.scale_model_input( + sample, t + ) + + else: + denoiser_input = sample + + # Predict noise level using denoiser using conditionings + pred = self.denoiser( + sample=denoiser_input, + timestep=t.to(z.device).repeat(denoiser_input.shape[0]), + conditioning=conditioning, + ) + + # Make one step on the reverse diffusion process + sample = self.sampling_noise_scheduler.step( + pred, t, sample, return_dict=False + )[0] + if i < len(self.sampling_noise_scheduler.timesteps) - 1: + timestep = ( + self.sampling_noise_scheduler.timesteps[i + 1] + .to(z.device) + .repeat(sample.shape[0]) + ) + sigmas = self._get_sigmas( + self.sampling_noise_scheduler, timestep, n_dim=4, device=z.device + ) + sample = sample + self.bridge_noise_sigma * ( + sigmas * (1.0 - sigmas) + ) ** 0.5 * torch.randn_like(sample) + sample = sample.to(z.dtype) + + if self.vae is not None: + decoded_sample = self.vae.decode(sample) + + else: + decoded_sample = sample + + return decoded_sample + + def log_samples( + self, + batch: Dict[str, Any], + input_shape: Optional[Tuple[int, int, int]] = None, + max_samples: Optional[int] = None, + num_steps: Union[int, List[int]] = 20, + ): + if isinstance(num_steps, int): + num_steps = [num_steps] + + logs = {} + + N = max_samples if max_samples is not None else len(batch[self.source_key]) + + batch = {k: v[:N] for k, v in batch.items()} + + # infer input shape based on VAE configuration if not passed + if input_shape is None: + if self.vae is not None: + # get input pixel size of the vae + input_shape = batch[self.target_key].shape[2:] + # rescale to latent size + input_shape = ( + self.vae.latent_channels, + input_shape[0] // self.vae.downsampling_factor, + input_shape[1] // self.vae.downsampling_factor, + ) + else: + raise ValueError( + "input_shape must be passed when no VAE is used in the model" + ) + + for num_step in num_steps: + source_image = batch[self.source_key] + source_image = torch.nn.functional.interpolate( + source_image, + size=batch[self.target_key].shape[2:], + mode="bilinear", + align_corners=False, + ).to(dtype=self.dtype) + if self.vae is not None: + z = self.vae.encode(source_image) + + else: + z = source_image + + with torch.autocast(dtype=self.dtype, device_type="cuda"): + logs[f"samples_{num_step}_steps"] = self.sample( + z, + num_steps=num_step, + conditioner_inputs=batch, + max_samples=N, + ) + + return logs diff --git a/scripts/lbm/tiler.py b/scripts/lbm/tiler.py new file mode 100644 index 000000000..529348332 --- /dev/null +++ b/scripts/lbm/tiler.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) diff --git a/scripts/lbm/unets/__init__.py b/scripts/lbm/unets/__init__.py new file mode 100644 index 000000000..178e63726 --- /dev/null +++ b/scripts/lbm/unets/__init__.py @@ -0,0 +1,15 @@ +""" +This module contains a collection of U-Net models. +The :mod:`cr.models.unets` module includes the following classes: + +- :class:`DiffusersUNet2DWrapper`: A 2D U-Net model for diffusers. +- :class:`DiffusersUNet2DCondWrapper`: A 2D U-Net model for diffusers with conditional input. +""" + +from .unet import DiffusersUNet2DCondWrapper, DiffusersUNet2DWrapper + + +__all__ = [ + "DiffusersUNet2DWrapper", + "DiffusersUNet2DCondWrapper", +] diff --git a/scripts/lbm/unets/unet.py b/scripts/lbm/unets/unet.py new file mode 100644 index 000000000..6aa6d7eee --- /dev/null +++ b/scripts/lbm/unets/unet.py @@ -0,0 +1,147 @@ +from typing import Dict, List, Optional, Union +import torch +from diffusers.models import UNet2DConditionModel, UNet2DModel + + +class DiffusersUNet2DWrapper(UNet2DModel): + """ + Wrapper for the UNet2DModel from diffusers + + See diffusers' UNet2DModel for more details + """ + + def __init__(self, *args, **kwargs): + UNet2DModel.__init__(self, *args, **kwargs) + + def forward( + self, + sample: torch.Tensor, + timestep: Union[torch.Tensor, float, int], + conditioning: Dict[str, torch.Tensor] = None, + *args, + **kwargs, + ): + """ + The forward pass of the model + + Args: + + sample (torch.Tensor): The input sample + timesteps (Union[torch.Tensor, float, int]): The number of timesteps + """ + if conditioning is not None: + class_labels = conditioning["cond"].get("vector", None) + concat = conditioning["cond"].get("concat", None) + + else: + class_labels = None + concat = None + + if concat is not None: + sample = torch.cat([sample, concat], dim=1) + + return super().forward(sample, timestep, class_labels).sample + + def freeze(self): + """ + Freeze the model + """ + self.eval() + for param in self.parameters(): + param.requires_grad = False + + +class DiffusersUNet2DCondWrapper(UNet2DConditionModel): + """ + Wrapper for the UNet2DConditionModel from diffusers + + See diffusers' Unet2DConditionModel for more details + """ + + def __init__(self, *args, **kwargs): + UNet2DConditionModel.__init__(self, *args, **kwargs) + # BaseModel.__init__(self, config=ModelConfig()) + + def forward( + self, + sample: torch.Tensor, + timestep: Union[torch.Tensor, float, int], + conditioning: Dict[str, torch.Tensor], + ip_adapter_cond_embedding: Optional[List[torch.Tensor]] = None, + down_block_additional_residuals: torch.Tensor = None, + mid_block_additional_residual: torch.Tensor = None, + down_intrablock_additional_residuals: torch.Tensor = None, + *args, + **kwargs, + ): + """ + The forward pass of the model + + Args: + + sample (torch.Tensor): The input sample + timesteps (Union[torch.Tensor, float, int]): The number of timesteps + conditioning (Dict[str, torch.Tensor]): The conditioning data + down_block_additional_residuals (List[torch.Tensor]): Residuals for the down blocks. + These residuals typically are used for the controlnet. + mid_block_additional_residual (List[torch.Tensor]): Residuals for the mid blocks. + These residuals typically are used for the controlnet. + down_intrablock_additional_residuals (List[torch.Tensor]): Residuals for the down intrablocks. + These residuals typically are used for the T2I adapters.middle block outputs. Defaults to False + """ + + assert isinstance(conditioning, dict), "conditionings must be a dictionary" + # assert "crossattn" in conditioning["cond"], "crossattn must be in conditionings" + + class_labels = conditioning["cond"].get("vector", None) + crossattn = conditioning["cond"].get("crossattn", None) + concat = conditioning["cond"].get("concat", None) + + # concat conditioning + if concat is not None: + sample = torch.cat([sample, concat], dim=1) + + # down_intrablock_additional_residuals needs to be cloned, since unet will modify it + if down_intrablock_additional_residuals is not None: + down_intrablock_additional_residuals_clone = [ + curr_residuals.clone() + for curr_residuals in down_intrablock_additional_residuals + ] + else: + down_intrablock_additional_residuals_clone = None + + # Check diffusers.models.embeddings.py > MultiIPAdapterImageProjectionLayer > forward() for implementation + # Exepected format : List[torch.Tensor] of shape (batch_size, num_image_embeds, embed_dim) + # with length = number of ip_adapters loaded in the ip_adapter_wrapper + if ip_adapter_cond_embedding is not None: + added_cond_kwargs = { + "image_embeds": [ + ip_adapter_embedding.unsqueeze(1) + for ip_adapter_embedding in ip_adapter_cond_embedding + ] + } + else: + added_cond_kwargs = None + + return ( + super() + .forward( + sample=sample, + timestep=timestep, + encoder_hidden_states=crossattn, + class_labels=class_labels, + added_cond_kwargs=added_cond_kwargs, + down_block_additional_residuals=down_block_additional_residuals, + mid_block_additional_residual=mid_block_additional_residual, + down_intrablock_additional_residuals=down_intrablock_additional_residuals_clone, + ) + .sample + ) + + def freeze(self): + """ + Freeze the model + """ + self.eval() + for param in self.parameters(): + param.requires_grad = False diff --git a/scripts/lbm/utils.py b/scripts/lbm/utils.py new file mode 100644 index 000000000..20c582ea9 --- /dev/null +++ b/scripts/lbm/utils.py @@ -0,0 +1,220 @@ +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 .embedders import ( + ConditionerWrapper, + LatentsConcatEmbedder, + LatentsConcatEmbedderConfig, +) +from .lbm import LBMConfig, LBMModel +from .unets import DiffusersUNet2DCondWrapper +from .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 diff --git a/scripts/lbm/vae/__init__.py b/scripts/lbm/vae/__init__.py new file mode 100644 index 000000000..05f8ab288 --- /dev/null +++ b/scripts/lbm/vae/__init__.py @@ -0,0 +1,5 @@ +from .autoencoderKL import AutoencoderKLDiffusers +from .autoencoderKL_config import AutoencoderKLDiffusersConfig + + +__all__ = ["AutoencoderKLDiffusers", "AutoencoderKLDiffusersConfig"] diff --git a/scripts/lbm/vae/autoencoderKL.py b/scripts/lbm/vae/autoencoderKL.py new file mode 100644 index 000000000..09a5523bb --- /dev/null +++ b/scripts/lbm/vae/autoencoderKL.py @@ -0,0 +1,135 @@ +import torch +from diffusers.models import AutoencoderKL +from ..base.base_model import BaseModel +from ..tiler 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 diff --git a/scripts/lbm/vae/autoencoderKL_config.py b/scripts/lbm/vae/autoencoderKL_config.py new file mode 100644 index 000000000..a2f1ed802 --- /dev/null +++ b/scripts/lbm/vae/autoencoderKL_config.py @@ -0,0 +1,25 @@ +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) diff --git a/scripts/lbm_ext.py b/scripts/lbm_ext.py new file mode 100644 index 000000000..67bcea1bf --- /dev/null +++ b/scripts/lbm_ext.py @@ -0,0 +1,136 @@ +from copy import deepcopy +from PIL import Image +import gradio as gr +from modules import scripts_manager, processing, shared, devices, sd_models + + +birefnet = None +model = None +model_type = '' +repos = { + 'Simple': None, + 'Normals': 'jasperai/LBM_normals', + 'Depth': 'jasperai/LBM_depth', + 'Relighting': 'jasperai/LBM_relighting', +} + +ASPECT_RATIOS = { + str(512 / 2048): (512, 2048), + str(1024 / 1024): (1024, 1024), + str(2048 / 512): (2048, 512), + str(896 / 1152): (896, 1152), + str(1152 / 896): (1152, 896), + str(512 / 1920): (512, 1920), + str(640 / 1536): (640, 1536), + str(768 / 1280): (768, 1280), + str(1280 / 768): (1280, 768), + str(1536 / 640): (1536, 640), + str(1920 / 512): (1920, 512), +} + + +class Script(scripts_manager.Script): + def title(self): + return 'LBM: Latent Bridge Matching' + + def show(self, is_img2img): + return is_img2img if shared.native else False + + # return signature is array of gradio components + def ui(self, _is_img2img): + with gr.Row(): + gr.HTML('  LBM: Latent Bridge Matching
') + with gr.Row(): + lbm_method = gr.Dropdown(label='LBM Method', choices=['Simple', 'Relighting', 'Normals', 'Depth'], value='Simple', elem_id='lbm_method') + with gr.Row(): + lbm_composite = gr.Checkbox(label='LBM Composite', value=True, elem_id='lbm_composite') + lbm_steps = gr.Slider(label='LBM Steps', minimum=1, maximum=20, step=1, value=1, elem_id='lbm_steps') + with gr.Row(): + bg_image = gr.Image(label='Background image', type='pil', height=512, elem_id='lbm_bg_image') + return [lbm_method, lbm_composite, lbm_steps, bg_image] + + def load(self, method: str): + global birefnet, model, model_type # pylint: disable=global-statement + import torch + if birefnet is None: + from transformers import AutoModelForImageSegmentation + birefnet = AutoModelForImageSegmentation.from_pretrained( + "ZhengPeng7/BiRefNet", + trust_remote_code=True, + torch_dtype=torch.float32, + ).to(dtype=torch.float32, device=devices.device) + if model is None or model_type != method: + repo_id = repos.get(method, None) + model_type = method + if repo_id is not None: + import huggingface_hub as hf + repo_file = hf.snapshot_download(repo_id, cache_dir=shared.opts.hfcache_dir) + from scripts.lbm import get_model + model = get_model( + repo_file, + save_dir=None, + torch_dtype=devices.dtype, + device=devices.device, + ).to(dtype=devices.dtype, device=devices.device) + + def run(self, p: processing.StableDiffusionProcessing, lbm_method, lbm_composite, lbm_steps, bg_image): # pylint: disable=arguments-differ, unused-argument + fg_image = getattr(p, 'init_images', None) + if fg_image is None or len(fg_image) == 0 or bg_image is None: + shared.log.error('LBM: no init images') + return None + else: + fg_image = fg_image[0] + + from installer import install + install('lpips') + + from torchvision.transforms import ToPILImage, ToTensor + from scripts.lbm import get_model, extract_object, resize_and_center_crop + + ori_h_bg, ori_w_bg = fg_image.size + ar_bg = ori_h_bg / ori_w_bg + closest_ar_bg = min(ASPECT_RATIOS, key=lambda x: abs(float(x) - ar_bg)) + dimensions_bg = ASPECT_RATIOS[closest_ar_bg] + + shared.log.info(f'LBM: method={lbm_method} steps={lbm_steps} size={dimensions_bg[0]}x{dimensions_bg[1]}') + self.load(lbm_method) + + if birefnet: + birefnet.to(device=devices.device) + if model: + model.to(device=devices.device) + + output_image = None + _, fg_mask = extract_object(birefnet, deepcopy(fg_image)) + fg_image = resize_and_center_crop(fg_image, dimensions_bg[0], dimensions_bg[1]) + fg_mask = resize_and_center_crop(fg_mask, dimensions_bg[0], dimensions_bg[1]) + bg_image = resize_and_center_crop(bg_image, dimensions_bg[0], dimensions_bg[1]) + img_pasted = Image.composite(fg_image, bg_image, fg_mask) + + if lbm_method == 'Simple': + output_image = img_pasted + else: + img_pasted_tensor = ToTensor()(img_pasted).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) * 2 - 1 + batch = { "source_image": img_pasted_tensor } + z_source = model.vae.encode(batch[model.source_key]) + output_image = model.sample( + z=z_source, + num_steps=lbm_steps, + conditioner_inputs=batch, + max_samples=1, + ) + output_image = (output_image[0].clamp(-1, 1).float().cpu() + 1) / 2 + output_image = ToPILImage()(output_image) + if lbm_composite: + output_image = Image.composite(output_image, bg_image, fg_mask) + + if birefnet: + birefnet.to(device=devices.cpu) + if model: + model.to(device=devices.cpu) + + if output_image is not None: + output_image.resize((ori_h_bg, ori_w_bg)) + return processing.Processed(p, [output_image]) + else: + return processing.Processed(p, []) diff --git a/scripts/mixture_of_diffusers.py b/scripts/mixture_of_diffusers.py index 0cbc60d4e..a3ee153ce 100644 --- a/scripts/mixture_of_diffusers.py +++ b/scripts/mixture_of_diffusers.py @@ -85,7 +85,7 @@ class Script(scripts_manager.Script): [x_tiles, y_tiles, x_overlap, y_overlap], prompts = args[:4], args[4:] if max(x_tiles, y_tiles) <= 1: return None - from scripts.mod import StableDiffusionXLTilingPipeline + from scripts.mod import StableDiffusionXLTilingPipeline # pylint: disable=no-name-in-module self.orig_pipe = shared.sd_model self.orig_attn = shared.opts.prompt_attention diff --git a/scripts/style_aligned_ext.py b/scripts/style_aligned_ext.py index 75e586971..a4e6cca54 100644 --- a/scripts/style_aligned_ext.py +++ b/scripts/style_aligned_ext.py @@ -64,7 +64,7 @@ class Script(scripts_manager.Script): shared.log.warning(f'SA: class={shared.sd_model.__class__.__name__} model={shared.sd_model_type} required={supported_model_list}') return None - from scripts.style_aligned import sa_handler, inversion + from scripts.style_aligned import sa_handler, inversion # pylint: disable=no-name-in-module handler = sa_handler.Handler(shared.sd_model) sa_args = sa_handler.StyleAlignedArgs(