Add refnet/models and ldm/models source files previously excluded by .gitignore

ldm/models/autoencoder.py

@@ -0,0 +1,37 @@
+import torch
+
+from ldm.modules.diffusionmodules.model import Encoder, Decoder
+from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
+
+
+class AutoencoderKL(torch.nn.Module):
+    def __init__(
+            self,
+            ddconfig,
+            embed_dim
+    ):
+        super().__init__()
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        assert ddconfig["double_z"]
+        self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        self.embed_dim = embed_dim
+
+    def encode(self, x):
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+        return posterior
+
+    def decode(self, z):
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        return dec
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+    @property
+    def dtype(self):
+        return self.decoder.conv_out.weight.dtype

refnet/models/basemodel.py

@@ -0,0 +1,439 @@
+import torch
+
+from refnet.util import exists, fitting_weights, instantiate_from_config, load_weights, delete_states
+from refnet.ldm import LatentDiffusion
+from typing import Union
+from refnet.sampling import (
+    UnetHook,
+    KDiffusionSampler,
+    DiffuserDenoiser,
+)
+
+
+
+class GuidanceFlag:
+    none = 0
+    reference = 1
+    sketch = 10
+    both = 11
+
+
+def reconstruct_cond(cond, uncond):
+    if not isinstance(uncond, list):
+        uncond = [uncond]
+    for k in cond.keys():
+        if k == "inpaint_bg":
+            continue
+        for uc in uncond:
+            if isinstance(cond[k], list):
+                cond[k] = [torch.cat([cond[k][i], uc[k][i]]) for i in range(len(cond[k]))]
+            elif isinstance(cond[k], torch.Tensor):
+                cond[k] = torch.cat([cond[k], uc[k]])
+    return cond
+
+
+class CustomizedLDM(LatentDiffusion):
+    def __init__(
+            self,
+            dtype = torch.float32,
+            sigma_max = None,
+            sigma_min = None,
+            *args,
+            **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.dtype = dtype
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+
+        self.model_list = {
+            "first": self.first_stage_model,
+            "cond": self.cond_stage_model,
+            "unet": self.model,
+        }
+        self.switch_cond_modules = ["cond"]
+        self.switch_main_modules = ["unet"]
+        self.retrieve_attn_modules()
+        self.retrieve_attn_layers()
+
+    def init_from_ckpt(
+            self,
+            path,
+            only_model = False,
+            logging = False,
+            make_it_fit = False,
+            ignore_keys: list[str] = (),
+    ):
+        sd = delete_states(load_weights(path), ignore_keys)
+        if make_it_fit:
+            sd = fitting_weights(self, sd)
+
+        missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model \
+            else self.model.load_state_dict(sd, strict=False)
+
+        filtered_missing = []
+        filtered_unexpect = []
+        for k in missing:
+            if not k.find("cond_stage_model") > -1 and not k.find("img_embedder") > -1 and not k.find("fg") > -1:
+                filtered_missing.append(k)
+        for k in unexpected:
+            if not k.find("cond_stage_model") > -1 and not k.find("img_embedder") > -1:
+                filtered_unexpect.append(k)
+
+        print(
+            f"Restored from {path} with {len(filtered_missing)} filtered missing and "
+            f"{len(filtered_unexpect)} filtered unexpected keys")
+        if logging:
+            if len(missing) > 0:
+                print(f"Filtered missing Keys: {filtered_missing}")
+            if len(unexpected) > 0:
+                print(f"Filtered unexpected Keys: {filtered_unexpect}")
+
+
+    def sample(
+            self,
+            cond: dict,
+            uncond: Union[dict, list[dict]] = None,
+            cfg_scale: Union[float, list[float]] = 1.,
+            bs: int = 1,
+            shape: Union[tuple, list] = None,
+            step: int = 20,
+            sampler = "DPM++ 3M SDE",
+            scheduler = "Automatic",
+            device = "cuda",
+            x_T = None,
+            seed = None,
+            deterministic = False,
+            **kwargs
+    ):
+        shape = shape or (self.channels, self.image_size, self.image_size)
+        x = x_T or torch.randn(bs, *shape, device=device)
+
+        if exists(uncond):
+            cond = reconstruct_cond(cond, uncond)
+
+        if sampler.startswith("diffuser"):
+            # Using huggingface diffuser noise sampler and scheduler
+            sampler = DiffuserDenoiser(
+                sampler,
+                prediction_type = "v_prediction" if self.parameterization == "v" else "epsilon",
+                use_karras = scheduler == "Karras"
+            )
+
+            samples = sampler(
+                x,
+                cond,
+                cond_scale=cfg_scale,
+                unet=self,
+                timesteps=step,
+                generator=torch.manual_seed(seed) if exists(seed) else None,
+                device=device
+            )
+
+        else:
+            # Using k-diffusion sampler and noise scheduler
+            seed = seed or torch.seed()
+            sampler = KDiffusionSampler(sampler, scheduler, self, device)
+            
+            sigmas = sampler.get_sigmas(step)
+            extra_args = {
+                "cond": cond,
+                "cond_scale": cfg_scale,
+            }
+            seed = [seed for _ in range(bs)] if deterministic else seed
+            samples = sampler(x, sigmas, extra_args, seed, deterministic, step)
+
+        return samples
+
+    def switch_to_fp16(self):
+        unet = self.model.diffusion_model
+        unet.input_blocks = unet.input_blocks.to(self.half_precision_dtype)
+        unet.middle_block = unet.middle_block.to(self.half_precision_dtype)
+        unet.output_blocks = unet.output_blocks.to(self.half_precision_dtype)
+        self.dtype = self.half_precision_dtype
+        unet.dtype = self.half_precision_dtype
+
+    def switch_to_fp32(self):
+        unet = self.model.diffusion_model
+        unet.input_blocks = unet.input_blocks.float()
+        unet.middle_block = unet.middle_block.float()
+        unet.output_blocks = unet.output_blocks.float()
+        self.dtype = torch.float32
+        unet.dtype = torch.float32
+
+    def switch_vae_to_fp16(self):
+        self.first_stage_model = self.first_stage_model.to(self.half_precision_dtype)
+
+    def switch_vae_to_fp32(self):
+        self.first_stage_model = self.first_stage_model.float()
+
+    def low_vram_shift(self, cuda_list: Union[str, list[str]]):
+        if not isinstance(cuda_list, list):
+            cuda_list = [cuda_list]
+
+        cpu_list = self.model_list.keys() - cuda_list
+        for model in cpu_list:
+            self.model_list[model] = self.model_list[model].cpu()
+        torch.cuda.empty_cache()
+
+        for model in cuda_list:
+            self.model_list[model] = self.model_list[model].cuda()
+
+
+    def retrieve_attn_modules(self):
+        from refnet.modules.transformer import BasicTransformerBlock
+        from refnet.sampling import torch_dfs
+
+        scale_factor_levels = {"high": 0.5, "low": 0.25, "bottom": 0.25}
+
+        attn_modules = []
+        for module in torch_dfs(self.model.diffusion_model):
+            if isinstance(module, BasicTransformerBlock):
+                attn_modules.append(module)
+
+        self.attn_modules = {
+            "high": [0, 1, 2, 3] + [64, 65, 66, 67, 68, 69],
+            "low": [i for i in range(4, 24)] + [i for i in range(34, 64)],
+            "bottom": [i for i in range(24, 34)],
+            "encoder": [i for i in range(24)],
+            "decoder": [i for i in range(34, len(attn_modules))]
+        }
+        self.attn_modules["modules"] = attn_modules
+
+        for k in ["high", "low", "bottom"]:
+            scale_factor = scale_factor_levels[k]
+            for attn in self.attn_modules[k]:
+                attn_modules[attn].scale_factor = scale_factor
+
+
+    def retrieve_attn_layers(self):
+        self.attn_layers = []
+        for module in (self.attn_modules["modules"]):
+            if hasattr(module, "attn2") and exists(getattr(module, "attn2")):
+                self.attn_layers.append(module.attn2)
+
+
+class CustomizedColorizer(CustomizedLDM):
+    def __init__(
+            self,
+            control_encoder_config,
+            proj_config,
+            token_type = "full",
+            *args,
+            **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.control_encoder = instantiate_from_config(control_encoder_config)
+        self.proj = instantiate_from_config(proj_config)
+        self.token_type = token_type
+        self.model_list.update({"control_encoder": self.control_encoder, "proj": self.proj})
+        self.switch_cond_modules += ["control_encoder", "proj"]
+
+
+    def switch_to_fp16(self):
+        self.control_encoder = self.control_encoder.to(self.half_precision_dtype)
+        super().switch_to_fp16()
+
+
+    def switch_to_fp32(self):
+        self.control_encoder = self.control_encoder.float()
+        super().switch_to_fp32()
+
+
+from refnet.modules.unet import hack_inference_forward
+class CustomizedWrapper:
+    def __init__(self):
+        self.scaling_sample = False
+        self.guidance_steps = (0, 1)
+        self.no_guidance_steps = (-0.05, 0.05)
+        hack_inference_forward(self.model.diffusion_model)
+
+    def adjust_reference_scale(self, scale_kwargs):
+        if isinstance(scale_kwargs, dict):
+            if scale_kwargs["level_control"]:
+                for key in scale_kwargs["scales"]:
+                    if key == "middle":
+                        continue
+                    for idx in self.attn_modules[key]:
+                        self.attn_modules["modules"][idx].reference_scale = scale_kwargs["scales"][key]
+            else:
+                for idx, s in enumerate(scale_kwargs["scales"]):
+                    self.attn_modules["modules"][idx].reference_scale = s
+        else:
+            for module in self.attn_modules["modules"]:
+                module.reference_scale = scale_kwargs
+
+    def adjust_fgbg_scale(self, fg_scale, bg_scale, merge_scale, mask_threshold):
+        for layer in self.attn_layers:
+            layer.fg_scale = fg_scale
+            layer.bg_scale = bg_scale
+            layer.merge_scale = merge_scale
+            layer.mask_threshold = mask_threshold
+        # for layer in self.attn_modules["modules"]:
+        #     layer.fg_scale = fg_scale
+        #     layer.bg_scale = bg_scale
+        #     layer.merge_scale = merge_scale
+        #     layer.mask_threshold = mask_threshold
+
+    def apply_model(self, x_noisy, t, cond):
+        tr = 1 - t[0] / (self.num_timesteps - 1)
+        crossattn = cond["context"][0]
+        if ((tr < self.guidance_steps[0] or tr > self.guidance_steps[1]) or
+                (tr >= self.no_guidance_steps[0] and tr <= self.no_guidance_steps[1])):
+            crossattn = torch.zeros_like(crossattn)[:, :1]
+        cond["context"] = [crossattn]
+        
+        model_cond = {k: v for k, v in cond.items() if k != "inpaint_bg"}
+        return self.model(x_noisy, t, **model_cond)
+
+
+    def prepare_conditions(self, *args, **kwargs):
+        raise NotImplementedError("Inputs preprocessing function is not implemented.")
+
+
+    def check_manipulate(self, scales):
+        if exists(scales) and len(scales) > 0:
+            for scale in scales:
+                if scale > 0:
+                    return True
+        return False
+
+    @torch.inference_mode()
+    def generate(
+            self,
+            # Conditional inputs
+            cond: dict,
+            ctl_scale: Union[float|list[float]],
+            merge_scale: float,
+            mask_scale: float,
+            mask_thresh: float,
+            mask_thresh_sketch: float,
+
+            # Sampling settings
+            sampler,
+            scheduler,
+            step: int,
+            bs: int,
+            gs: list[float],
+            strength: Union[float, list[float]],
+            fg_strength: float,
+            bg_strength: float,
+            seed: int,
+            start_step: float = 0.0,
+            end_step: float = 1.0,
+            no_start_step: float = -0.05,
+            no_end_step: float = -0.05,
+            deterministic: bool = False,
+            style_enhance: bool = False,
+            bg_enhance: bool = False,
+            fg_enhance: bool = False,
+            latent_inpaint: bool = False,
+            height: int = 512,
+            width: int = 512,
+
+            # Injection settings
+            injection: bool = False,
+            injection_cfg: float = 0.5,
+            injection_control: float = 0,
+            injection_start_step: float = 0,
+            hook_xr: torch.Tensor = None,
+            hook_xs: torch.Tensor = None,
+
+            # Additional settings
+            low_vram: bool = True,
+            return_intermediate = False,
+            manipulation_params = None,
+            **kwargs,
+    ):
+        """
+            User interface function.
+        """
+        hook_unet = UnetHook()
+
+        self.guidance_steps = (start_step, end_step)
+        self.no_guidance_steps = (no_start_step, no_end_step)
+        self.adjust_reference_scale(strength)
+        self.adjust_fgbg_scale(fg_strength, bg_strength, merge_scale, mask_thresh_sketch)
+
+        if low_vram:
+            self.low_vram_shift(self.switch_cond_modules)
+        else:
+            self.low_vram_shift(list(self.model_list.keys()))
+
+        c, uc = self.prepare_conditions(
+            bs = bs,
+            control_scale = ctl_scale,
+            merge_scale = merge_scale,
+            mask_scale = mask_scale,
+            mask_threshold_ref = mask_thresh,
+            mask_threshold_sketch = mask_thresh_sketch,
+            style_enhance = style_enhance,
+            bg_enhance = bg_enhance,
+            fg_enhance = fg_enhance,
+            latent_inpaint = latent_inpaint,
+            height = height,
+            width = width,
+            bg_strength = bg_strength,
+            low_vram = low_vram,
+            **cond,
+            **manipulation_params,
+            **kwargs
+        )
+
+        cfg = int(gs[0] > 1) * GuidanceFlag.reference + int(gs[1] > 1) * GuidanceFlag.sketch
+        gr_indice = [] if (cfg == GuidanceFlag.none or cfg == GuidanceFlag.sketch) else [i for i in range(bs, bs*2)]
+        repeat = 1
+        if cfg == GuidanceFlag.none:
+            gs = 1
+            uc = None
+        if cfg == GuidanceFlag.reference:
+            gs = gs[0]
+            uc = uc[0]
+            repeat = 2
+        if cfg == GuidanceFlag.sketch:
+            gs = gs[1]
+            uc = uc[1]
+            repeat = 2
+        if cfg == GuidanceFlag.both:
+            repeat = 3
+
+        if low_vram:
+            self.low_vram_shift("first")
+
+        if injection:
+            rx = self.get_first_stage_encoding(hook_xr.to(self.first_stage_model.dtype))
+            hook_unet.enhance_reference(
+                model = self.model,
+                ldm = self,
+                bs = bs * repeat,
+                s = -hook_xr.to(self.dtype),
+                r = rx,
+                style_cfg = injection_cfg,
+                control_cfg = injection_control,
+                gr_indice = gr_indice,
+                start_step = injection_start_step,
+            )
+
+        if low_vram:
+            self.low_vram_shift(self.switch_main_modules)
+
+        z = self.sample(
+            cond = c,
+            uncond = uc,
+            bs = bs,
+            shape = (self.channels, height // 8, width // 8),
+            cfg_scale = gs,
+            step = step,
+            sampler = sampler,
+            scheduler = scheduler,
+            seed = seed,
+            deterministic = deterministic,
+            return_intermediate = return_intermediate,
+        )
+
+        if injection:
+            hook_unet.restore(self.model)
+
+        if low_vram:
+            self.low_vram_shift("first")
+        return self.decode_first_stage(z.to(self.first_stage_model.dtype))

refnet/models/colorizerXL.py

@@ -0,0 +1,201 @@
+
+import torch
+import torch.nn.functional as F
+
+from ..modules.reference_net import hack_inference_forward
+from ..models.basemodel import CustomizedColorizer, CustomizedWrapper
+from ..modules.lora import LoraModules
+from ..util import exists, expand_to_batch_size, instantiate_from_config, get_crop_scale, resize_and_crop
+
+
+
+class InferenceWrapper(CustomizedWrapper, CustomizedColorizer):
+    def __init__(
+            self,
+            scalar_embedder_config,
+            img_embedder_config,
+            lora_config = None,
+            logits_embed = False,
+            *args,
+            **kwargs
+    ):
+        CustomizedColorizer.__init__(self, version="sdxl", *args, **kwargs)
+        CustomizedWrapper.__init__(self)
+
+        self.scalar_embedder = instantiate_from_config(scalar_embedder_config)
+        self.img_embedder = instantiate_from_config(img_embedder_config)
+        self.loras = LoraModules(self, **lora_config) if exists(lora_config) else None
+        self.logits_embed = logits_embed
+
+        new_model_list = {
+            "scalar_embedder": self.scalar_embedder,
+            "img_embedder": self.img_embedder,
+            # "style_encoder": self.style_encoder,
+        }
+        self.switch_cond_modules += list(new_model_list.keys())
+        self.model_list.update(new_model_list)
+
+    def retrieve_attn_modules(self):
+        scale_factor_levels = {"high": 0.5, "low": 0.25, "bottom": 0.25}
+
+        from refnet.modules.transformer import BasicTransformerBlock
+        from refnet.sampling import torch_dfs
+
+        attn_modules = []
+        for module in torch_dfs(self.model.diffusion_model):
+            if isinstance(module, BasicTransformerBlock):
+                attn_modules.append(module)
+
+        self.attn_modules = {
+            "high": [0, 1, 2, 3] + [64, 65, 66, 67, 68, 69],
+            "low": [i for i in range(4, 24)] + [i for i in range(34, 64)],
+            "bottom": [i for i in range(24, 34)],
+            "encoder": [i for i in range(24)],
+            "decoder": [i for i in range(34, len(attn_modules))]
+        }
+        self.attn_modules["modules"] = attn_modules
+
+        for k in ["high", "low", "bottom"]:
+            scale_factor = scale_factor_levels[k]
+            for attn in self.attn_modules[k]:
+                attn_modules[attn].scale_factor = scale_factor
+
+    def adjust_reference_scale(self, scale_kwargs):
+        for module in self.attn_modules["modules"]:
+            module.reference_scale = scale_kwargs["scales"]["encoder"]
+
+    def adjust_masked_attn(self, scale, mask_threshold, merge_scale):
+        for layer in self.attn_layers:
+            layer.mask_scale = scale
+            layer.mask_threshold = mask_threshold
+            layer.merge_scale = merge_scale
+
+    def rescale_size(self, x: torch.Tensor, height, width):
+        oh, ow = x.shape[2:]
+        if oh < height or ow < width:
+            dh, dw = height - oh, width - ow
+            if dh > dw:
+                iw = ow + int(dh * ow/oh)
+                ih = height
+            else:
+                ih = oh + int(dw * oh/ow)
+                iw = width
+        else:
+            ih, iw = oh, ow
+        return torch.Tensor([ih]), torch.Tensor([iw])
+
+    def get_learned_embedding(self, c, bg=False, mapping=False, sketch=None, *args, **kwargs):
+        clip_emb = self.cond_stage_model.encode(c, "full").detach()
+        wd_emb, logits = self.img_embedder.encode(c, pooled=False, return_logits=True)
+        cls_emb, local_emb = clip_emb[:, :1], clip_emb[:, 1:]
+
+        if mapping:
+            _, sketch_logits = self.img_embedder.encode(-sketch, pooled=False, return_logits=True)
+            sketch_logits.mean(dim=1, keepdim=True)
+            logits = self.img_embedder.geometry_update(logits, sketch_logits)
+        emb = self.proj(clip_emb, logits if self.logits_embed else wd_emb, bg)
+        return emb, cls_emb
+
+    def prepare_conditions(
+            self,
+            bs,
+            sketch,
+            reference,
+            height,
+            width,
+            control_scale = (1., 1., 1., 1.),
+            merge_scale = 0,
+            mask_scale = 1.,
+            fg_scale = 1.,
+            bg_scale = 1.,
+            smask = None,
+            rmask = None,
+            mask_threshold_ref = 0.,
+            mask_threshold_sketch = 0.,
+            style_enhance = False,
+            fg_enhance = False,
+            bg_enhance = False,
+            background = None,
+            targets = None,
+            anchors = None,
+            controls = None,
+            target_scales = None,
+            enhances = None,
+            thresholds_list = None,
+            geometry_map = False,
+            latent_inpaint = False,
+            low_vram = False,
+            *args,
+            **kwargs
+    ):
+        # prepare reference embedding
+        # manipulate = self.check_manipulate(target_scales)
+        c = {}
+        uc = [{}, {}]
+
+        if exists(reference):
+            emb, cls_emb = self.get_learned_embedding(reference, sketch=sketch, mapping=geometry_map)
+        else:
+            emb, cls_emb = map(lambda t: torch.zeros_like(t), self.get_learned_embedding(sketch))
+
+        h, w, score = torch.Tensor([height]), torch.Tensor([width]), torch.Tensor([7.])
+        y = torch.cat(self.scalar_embedder(torch.cat([(h*w)**0.5, score])).cuda().chunk(2), 1)
+
+        if bg_enhance:
+            assert exists(rmask) and exists(smask)
+
+            if low_vram:
+                self.low_vram_shift(["first", "cond", "img_embedder", "proj"])
+            
+            if latent_inpaint and exists(background):
+                bgh, bgw = background.shape[2:]
+                ch, cw = get_crop_scale(torch.tensor([height]), torch.tensor([width]), bgh, bgw)
+                hs_bg = self.get_first_stage_encoding(resize_and_crop(background, ch, cw, height, width).to(self.first_stage_model.dtype))
+                bg_emb, _ = self.get_learned_embedding(background, bg=True)
+                hs_bg = expand_to_batch_size(hs_bg, bs)
+                c.update({"inpaint_bg": hs_bg})
+            else:
+                if exists(background):
+                    bg_emb, _ = self.get_learned_embedding(background, bg=True)
+                else:
+                    bg_emb, _ = self.get_learned_embedding(
+                        torch.where(rmask < mask_threshold_ref, reference, torch.ones_like(reference)),
+                        True
+                    )
+            emb = torch.cat([emb, bg_emb], 1)
+
+        if fg_enhance and exists(self.loras):
+            self.loras.switch_lora(True, "foreground")
+            if not bg_enhance:
+                emb = emb.repeat(1, 2, 1)
+
+        if fg_enhance or bg_enhance:
+            # sketch mask for cross-attention
+            smask = expand_to_batch_size(smask.to(self.dtype), bs)
+            for d in [c] + uc:
+                d.update({"mask": F.interpolate(smask, scale_factor=0.125)})
+        elif exists(self.loras):
+            self.loras.switch_lora(False)
+
+        sketch = sketch.to(self.dtype)
+        context = expand_to_batch_size(emb, bs).to(self.dtype)
+        y = expand_to_batch_size(y, bs)
+        uc_context = torch.zeros_like(context)
+
+        control = []
+        uc_control = []
+        if low_vram:
+            self.low_vram_shift(["control_encoder"])
+        encoded_sketch = self.control_encoder(
+            torch.cat([sketch, -torch.ones_like(sketch)], 0)
+        )
+        for idx, es in enumerate(encoded_sketch):
+            es = es * control_scale[idx]
+            ec, uec = es.chunk(2)
+            control.append(expand_to_batch_size(ec, bs))
+            uc_control.append(expand_to_batch_size(uec, bs))
+
+        c.update({"control": control, "context": [context], "y": [y]})
+        uc[0].update({"control": control, "context": [uc_context], "y": [y]})
+        uc[1].update({"control": uc_control, "context": [context], "y": [y]})
+        return c, uc

refnet/models/v2-colorizerXL.py

@@ -0,0 +1,386 @@
+from refnet.models.basemodel import CustomizedColorizer, CustomizedWrapper
+from refnet.util import *
+from refnet.modules.lora import LoraModules
+from refnet.modules.reference_net import hack_unet_forward, hack_inference_forward
+from refnet.sampling.hook import ReferenceAttentionControl
+
+
+class InferenceWrapperXL(CustomizedWrapper, CustomizedColorizer):
+    def __init__(
+            self,
+            scalar_embedder_config,
+            img_embedder_config,
+            fg_encoder_config = None,
+            bg_encoder_config = None,
+            style_encoder_config = None,
+            lora_config = None,
+            logits_embed = False,
+            controller = False,
+            *args,
+            **kwargs
+    ):
+        CustomizedColorizer.__init__(self, version="sdxl", *args, **kwargs)
+        CustomizedWrapper.__init__(self)
+
+        self.logits_embed = logits_embed
+
+        (
+            self.scalar_embedder,
+            self.img_embedder,
+            self.fg_encoder,
+            self.bg_encoder,
+            self.style_encoder
+        ) = map(
+            lambda t: instantiate_from_config(t) if exists(t) else None,
+            (
+                scalar_embedder_config,
+                img_embedder_config,
+                fg_encoder_config,
+                bg_encoder_config,
+                style_encoder_config
+            )
+        )
+        self.loras = LoraModules(self, **lora_config)
+
+        if controller:
+            self.controller = ReferenceAttentionControl(
+                # time_embed_ch = self.model.diffusion_model.model_channels * 4,
+                reader_module = self.model.diffusion_model,
+                writer_module = self.bg_encoder,
+                # only_decoder = True
+            )
+        else:
+            self.controller = None
+
+        new_model_list = {
+            # "style_encoder": self.style_encoder,
+            "scalar_embedder": self.scalar_embedder,
+            "img_embedder": self.img_embedder,
+            # "controller": self.controller
+        }
+
+        hack_unet_forward(self.model.diffusion_model)
+        if exists(self.fg_encoder):
+            hack_inference_forward(self.fg_encoder)
+            new_model_list["fg_encoder"] = self.fg_encoder
+        if exists(self.bg_encoder):
+            hack_inference_forward(self.bg_encoder)
+            new_model_list["bg_encoder"] = self.bg_encoder
+        # hack_inference_forward(self.bg_encoder)
+        # hack_inference_forward(self.style_encoder)
+
+        self.switch_cond_modules += list(new_model_list.keys())
+        # self.switch_main_modules += ["controller"]
+        self.model_list.update(new_model_list)
+
+
+    def switch_to_fp16(self):
+        super().switch_to_fp16()
+        self.model.diffusion_model.map_modules.to(self.half_precision_dtype)
+        self.model.diffusion_model.warp_modules.to(self.half_precision_dtype)
+        self.model.diffusion_model.style_modules.to(self.half_precision_dtype)
+        self.model.diffusion_model.conv_fg.to(self.half_precision_dtype)
+
+        if exists(self.fg_encoder):
+            self.fg_encoder.to(self.half_precision_dtype)
+            self.fg_encoder.dtype = self.half_precision_dtype
+            self.fg_encoder.time_embed.float()
+        if exists(self.bg_encoder):
+            self.bg_encoder.to(self.half_precision_dtype)
+            self.bg_encoder.dtype = self.half_precision_dtype
+            self.bg_encoder.time_embed.float()
+        # self.style_encoder.to(self.half_precision_dtype)
+        # self.style_encoder.dtype = self.half_precision_dtype
+        # self.style_encoder.time_embed.float()
+
+    def switch_to_fp32(self):
+        super().switch_to_fp32()
+        self.model.diffusion_model.map_modules.float()
+        self.model.diffusion_model.warp_modules.float()
+        self.model.diffusion_model.style_modules.float()
+
+        self.fg_encoder.float()
+        self.bg_encoder.float()
+        # self.style_encoder.float()
+
+        self.fg_encoder.dtype = torch.float32
+        self.bg_encoder.dtype = torch.float32
+        # self.style_encoder.dtype = torch.float32
+
+    def rescale_size(self, x: torch.Tensor, height, width):
+        oh, ow = x.shape[2:]
+        if oh < height or ow < width:
+            dh, dw = height - oh, width - ow
+            if dh > dw:
+                iw = ow + int(dh * ow/oh)
+                ih = height
+            else:
+                ih = oh + int(dw * oh/ow)
+                iw = width
+        else:
+            ih, iw = oh, ow
+        return torch.tensor([ih]), torch.tensor([iw])
+
+    def rescale_background_size(self, x, height, width):
+        oh, ow = x.shape[2:]
+        if oh < height or ow < width:
+            # A simple bias to avoid deterioration caused by reference resolution
+            mind = max(height, width)
+            ih = oh + mind
+            iw = ow / oh * ih
+        else:
+            ih, iw = oh, ow
+        # rh, rw = ih / height, iw / width
+        return torch.tensor([ih]), torch.tensor([iw])
+
+    def get_learned_embedding(self, c, bg=False, sketch=None, mapping=False, *args, **kwargs):
+        clip_emb = self.cond_stage_model.encode(c, "full").detach()
+        wd_emb, logits = self.img_embedder.encode(c, pooled=False, return_logits=True)
+        cls_emb, local_emb = clip_emb[:, :1], clip_emb[:, 1:]
+
+        if self.logits_embed and exists(sketch) and mapping:
+            _, sketch_logits = self.img_embedder.encode(-sketch, pooled=True, return_logits=True)
+            logits = self.img_embedder.geometry_update(logits, sketch_logits)
+
+        if self.logits_embed:
+            emb = self.proj(clip_emb, logits, bg)[0]
+        else:
+            emb = self.proj(clip_emb, wd_emb, bg)
+        return emb.to(self.dtype), cls_emb.to(self.dtype)
+
+    def prepare_conditions(
+            self,
+            bs,
+            sketch,
+            reference,
+            height,
+            width,
+            control_scale = 1,
+            mask_scale = 1,
+            merge_scale = 0.,
+            cond_aug = 0.,
+            background = None,
+            smask = None,
+            rmask = None,
+            mask_threshold_ref = 0.,
+            mask_threshold_sketch = 0.,
+            style_enhance = False,
+            fg_enhance = False,
+            bg_enhance = False,
+            latent_inpaint = False,
+            fg_disentangle_scale = 1.,
+            targets = None,
+            anchors = None,
+            controls = None,
+            target_scales = None,
+            enhances = None,
+            thresholds_list = None,
+            low_vram = False,
+            *args,
+            **kwargs
+    ):
+        def prepare_style_modulations(y):
+            # Style enhancement part
+            z_ref = self.get_first_stage_encoding(warp_resize(reference, (height, width)))
+            if exists(background) and merge_scale > 0:
+                rh, rw = self.rescale_size(background, height, width)
+                z_bg = self.get_first_stage_encoding(warp_resize(background, (height, width)))
+                bg_emb, bg_cls_emb = self.get_learned_embedding(background)
+                scalar_embed = torch.cat(
+                    self.scalar_embedder(torch.cat([rh, rw, ct, cl, h, w])).chunk(6), 1
+                ).to(bg_emb.device)
+                bgy = torch.cat([bg_cls_emb.squeeze(1), scalar_embed], 1).to(self.dtype)
+
+                style_modulations = self.style_encoder(
+                    torch.cat([z_ref, z_bg]),
+                    timesteps = torch.zeros((2,), dtype=torch.long, device=z_ref.device),
+                    context = torch.cat([emb, bg_emb]),
+                    y = torch.cat([y, bgy])
+                )
+
+                for idx, m in enumerate(style_modulations):
+                    fg, bg = m.chunk(2)
+                    m = fg * (1-merge_scale) + merge_scale * bg
+                    style_modulations[idx] = expand_to_batch_size(m, bs).to(self.dtype)
+
+            else:
+                z_bg = None
+                bg_emb = None
+                bgy = None
+                style_modulations = self.style_encoder(
+                    z_ref,
+                    timesteps = torch.zeros((1,), dtype=torch.long, device=z_ref.device),
+                    context = emb,
+                    y = y,
+                )
+                style_modulations = [expand_to_batch_size(m, bs).to(self.dtype) for m in style_modulations]
+
+            return style_modulations, z_bg, bg_emb, bgy
+
+
+        def prepare_background_latents(z_bg, bg_emb, bgy):
+            # Background enhancement part
+            bgh, bgw = background.shape[2:] if exists(background) else reference.shape[2:]
+            ch, cw = get_crop_scale(h, w, bgh, bgw)
+
+            if low_vram:
+                self.low_vram_shift(["first", "cond", "img_embedder"])
+            if latent_inpaint and exists(background):
+                hs_bg = self.get_first_stage_encoding(resize_and_crop(background, ch, cw, height, width))
+                bg_emb, cls_emb = self.get_learned_embedding(background)
+
+            else:
+                if not exists(z_bg):
+                    bgy = torch.cat(
+                        self.scalar_embedder(torch.tensor([ct, cl, ch, cw])).chunk(4), 1
+                        # self.scalar_embedder(torch.tensor([bgh / bgw, h / w, ct, cl, ch, cw])).chunk(6), 1
+                    ).to(self.dtype).cuda()
+
+                    if exists(background):
+                        # bgh, bgw = self.rescale_background_size(background, height, width)
+                        z_bg = self.get_first_stage_encoding(warp_resize(background, (height, width)))
+                        bg_emb, cls_emb = self.get_learned_embedding(background)
+                        # scalar_embed = torch.cat(self.scalar_embedder(torch.cat([bgh, bgw, ct, cl, h, w])).chunk(6), 1).cuda()
+                        # bgy = torch.cat([cls_emb.squeeze(1), scalar_embed], 1).to(self.dtype)
+                    else:
+                        xbg = torch.where(rmask < mask_threshold_ref, reference, torch.ones_like(reference))
+                        z_bg = self.get_first_stage_encoding(warp_resize(xbg, (height, width)))
+                        bg_emb, cls_emb = self.get_learned_embedding(xbg)
+
+                if low_vram:
+                    self.low_vram_shift(["bg_encoder"])
+                hs_bg = self.bg_encoder(
+                    x = torch.cat([
+                        z_bg,
+                        # torch.where(
+                        #     smask > mask_threshold_sketch,
+                        #     torch.zeros_like(smask),
+                        #     F.interpolate(warp_resize(rmask, (height, width)), scale_factor=0.125)
+                        # )
+                        F.interpolate(warp_resize(smask, (height, width)), scale_factor=0.125),
+                        F.interpolate(warp_resize(rmask, (height, width)), scale_factor=0.125)
+                    ], 1),
+                    timesteps = torch.zeros((1,), dtype=torch.long, device=z_bg.device),
+                    # context = bg_emb,
+                    y = bgy.to(self.dtype),
+                )
+            return hs_bg, bg_emb
+
+        self.loras.recover_lora()
+        # prepare reference embedding
+        # manipulate = self.check_manipulate(target_scales)
+        c = {}
+        uc = [{}, {}]
+        self.loras.switch_lora(False)
+        # self.loras.recover_lora()
+
+        if exists(reference):
+            emb, cls_emb = self.get_learned_embedding(reference, sketch=sketch)
+            # rh, rw = reference.shape[2:]
+            # rh, rw = self.rescale_background_size(reference, height, width)
+
+        else:
+            emb, cls_emb = map(lambda t: torch.zeros_like(t), self.get_learned_embedding(sketch))
+            # rh, rw = torch.Tensor([height]), torch.Tensor([width])
+
+        ct, cl = torch.Tensor([0]), torch.Tensor([0])
+        # h, w = torch.Tensor([height]), torch.Tensor([width])
+        # scalar_embed = torch.cat(self.scalar_embedder(torch.cat([rh, rw, ct, cl, h, w])).chunk(6), 1).cuda()
+        # y = torch.cat([cls_emb.squeeze(1), scalar_embed], 1)
+        # y = self.scalar_embedder((h*w)**0.5).cuda()
+        # y = torch.cat(self.scalar_embedder(torch.cat([h, w])).chunk(2), 1).cuda()
+        h, w, score = torch.Tensor([height]), torch.Tensor([width]), torch.Tensor([7.])
+        y = torch.cat(self.scalar_embedder(torch.cat([(h * w) ** 0.5, score])).cuda().chunk(2), 1)
+
+        z_bg, bg_emb, bgy = None, None, None
+
+        # Style enhance part
+        if style_enhance:
+            style_modulations, z_bg, bg_emb, bgy = prepare_style_modulations(y)
+            for d in [c] + uc:
+                d.update({"style_modulations": style_modulations})
+
+        # Foreground enhance part
+        if fg_enhance:
+            assert exists(smask) and exists(rmask)
+            self.loras.switch_lora(True, "foreground")
+            if low_vram:
+                self.low_vram_shift(["first"])
+            z_fg = self.get_first_stage_encoding(warp_resize(
+                torch.where(rmask >= mask_threshold_ref, reference, torch.ones_like(reference)),
+                (height, width)
+            )) * fg_disentangle_scale
+            # z_ref = default(z_ref, self.get_first_stage_encoding(warp_resize(reference, (height, width))))
+            # self.loras.switch_lora(True, False)
+            self.loras.adjust_lora_scales(fg_disentangle_scale, "foreground")
+            if low_vram:
+                self.low_vram_shift(["fg_encoder"])
+            hs_fg = self.fg_encoder(
+                z_fg,
+                timesteps = torch.zeros((1,), dtype=torch.long, device=z_fg.device),
+            )
+            # hs_fg = [hs * fg_disentangle_scale for hs in hs_fg]
+            hs_fg = expand_to_batch_size(hs_fg, bs)
+            for d in [c] + uc:
+                d.update({
+                    "hs_fg": hs_fg,
+                    "inject_mask": expand_to_batch_size(smask, bs),
+                })
+            # for d in [c] + uc:
+                # d.update({"z_fg": expand_to_batch_size(z_fg, bs)})
+
+        # Background enhance part
+        if bg_enhance:
+            assert exists(rmask) and exists(smask)
+            # if not self.controller.hooked:
+            #     self.controller.register("read", self.model.diffusion_model)
+            # self.loras.switch_lora(False, True)
+            hs_bg, bg_emb = prepare_background_latents(z_bg, bg_emb, default(bgy, y))
+            self.loras.switch_lora(True, "background")
+            if latent_inpaint and exists(background):
+                hs_bg = expand_to_batch_size(hs_bg, bs)
+                c.update({"inpaint_bg": hs_bg})
+            elif exists(self.controller):
+                # self.loras.merge_lora()
+                self.controller.update()
+            else:
+                hs_bg = expand_to_batch_size(hs_bg, bs)
+                for d in [c] + uc:
+                    d.update({"hs_bg": hs_bg})
+
+        elif exists(self.controller):
+            # self.controller.reader_restore()
+            self.controller.clean()
+
+        if fg_enhance or bg_enhance:
+            # need to activate mask-guided split cross-attetnion
+            emb = torch.cat([emb, default(bg_emb, emb)], 1)
+            smask = expand_to_batch_size(smask.to(self.dtype), bs)
+            for d in [c] + uc:
+                d.update({"mask": F.interpolate(smask, scale_factor=0.125), "threshold": mask_threshold_sketch})
+
+        # if fg_enhance and bg_enhance:
+            # self.loras.switch_lora(True, True)
+        sketch = sketch.to(self.dtype)
+        context = expand_to_batch_size(emb, bs).to(self.dtype)
+        y = expand_to_batch_size(y, bs).float()
+        uc_context = torch.zeros_like(context)
+
+        control = []
+        uc_control = []
+        if low_vram:
+            self.low_vram_shift(["control_encoder"])
+        encoded_sketch = self.control_encoder(
+            torch.cat([sketch, -torch.ones_like(sketch)], 0)
+        )
+        for idx, es in enumerate(encoded_sketch):
+            es = es * control_scale[idx]
+            ec, uec = es.chunk(2)
+            control.append(expand_to_batch_size(ec, bs))
+            uc_control.append(expand_to_batch_size(uec, bs))
+
+        self.loras.merge_lora()
+        c.update({"control": control, "context": [context], "y": [y]})
+        uc[0].update({"control": control, "context": [uc_context], "y": [y]})
+        uc[1].update({"control": uc_control, "context": [context], "y": [y]})
+        return c, uc