libcom 0.1.8__py3-none-any.whl

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/util.py

@@ -0,0 +1,200 @@
+import importlib
+import math
+
+import cv2
+import torch
+import numpy as np
+
+import os
+from safetensors.torch import load_file
+
+from inspect import isfunction
+from PIL import Image, ImageDraw, ImageFont
+
+
+def log_txt_as_img(wh, xc, size=10):
+    # wh a tuple of (width, height)
+    # xc a list of captions to plot
+    b = len(xc)
+    txts = list()
+    for bi in range(b):
+        txt = Image.new("RGB", wh, color="white")
+        draw = ImageDraw.Draw(txt)
+        font = ImageFont.truetype('assets/DejaVuSans.ttf', size=size)
+        nc = int(40 * (wh[0] / 256))
+        lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
+
+        try:
+            draw.text((0, 0), lines, fill="black", font=font)
+        except UnicodeEncodeError:
+            print("Cant encode string for logging. Skipping.")
+
+        txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
+        txts.append(txt)
+    txts = np.stack(txts)
+    txts = torch.tensor(txts)
+    return txts
+
+
+def ismap(x):
+    if not isinstance(x, torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] > 3)
+
+
+def isimage(x):
+    if not isinstance(x, torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def mean_flat(tensor):
+    """
+    https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def count_params(model, verbose=False):
+    total_params = sum(p.numel() for p in model.parameters())
+    if verbose:
+        print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
+    return total_params
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        if config == '__is_first_stage__':
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+checkpoint_dict_replacements = {
+    'cond_stage_model.transformer.text_model.embeddings.': 'cond_stage_model.transformer.embeddings.',
+    'cond_stage_model.transformer.text_model.encoder.': 'cond_stage_model.transformer.encoder.',
+    'cond_stage_model.transformer.text_model.final_layer_norm.': 'cond_stage_model.transformer.final_layer_norm.',
+}
+
+
+def transform_checkpoint_dict_key(k):
+    for text, replacement in checkpoint_dict_replacements.items():
+        if k.startswith(text):
+            k = replacement + k[len(text):]
+
+    return k
+
+
+def get_state_dict_from_checkpoint(pl_sd):
+    pl_sd = pl_sd.pop("state_dict", pl_sd)
+    pl_sd.pop("state_dict", None)
+
+    sd = {}
+    for k, v in pl_sd.items():
+        new_key = transform_checkpoint_dict_key(k)
+
+        if new_key is not None:
+            sd[new_key] = v
+
+    pl_sd.clear()
+    pl_sd.update(sd)
+
+    return pl_sd
+
+
+def read_state_dict(checkpoint_file, print_global_state=False):
+    _, extension = os.path.splitext(checkpoint_file)
+    if extension.lower() == ".safetensors":
+        pl_sd = load_file(checkpoint_file, device='cpu')
+    else:
+        pl_sd = torch.load(checkpoint_file, map_location='cpu')
+
+    if print_global_state and "global_step" in pl_sd:
+        print(f"Global Step: {pl_sd['global_step']}")
+
+    sd = get_state_dict_from_checkpoint(pl_sd)
+    return sd
+
+
+def load_model_from_config(config, ckpt, vae_ckpt=None, verbose=False):
+    print(f"Loading model from {ckpt}")
+    sd = read_state_dict(ckpt)
+    model = instantiate_from_config(config.model)
+    m, u = model.load_state_dict(sd, strict=False)
+    if len(m) > 0 and verbose:
+        print("missing keys:")
+        print(m)
+    if len(u) > 0 and verbose:
+        print("unexpected keys:")
+        print(u)
+
+    if 'anything' in ckpt.lower() and vae_ckpt is None:
+        vae_ckpt = 'models/anything-v4.0.vae.pt'
+
+    if vae_ckpt is not None and vae_ckpt != 'None':
+        print(f"Loading vae model from {vae_ckpt}")
+        vae_sd = torch.load(vae_ckpt, map_location="cpu")
+        if "global_step" in vae_sd:
+            print(f"Global Step: {vae_sd['global_step']}")
+        sd = vae_sd["state_dict"]
+        m, u = model.first_stage_model.load_state_dict(sd, strict=False)
+        if len(m) > 0 and verbose:
+            print("missing keys:")
+            print(m)
+        if len(u) > 0 and verbose:
+            print("unexpected keys:")
+            print(u)
+
+    model.cuda()
+    model.eval()
+    return model
+
+
+def resize_numpy_image(image, max_resolution=512 * 512, resize_short_edge=None):
+    h, w = image.shape[:2]
+    if resize_short_edge is not None:
+        k = resize_short_edge / min(h, w)
+    else:
+        k = max_resolution / (h * w)
+        k = k**0.5
+    h = int(np.round(h * k / 64)) * 64
+    w = int(np.round(w * k / 64)) * 64
+    image = cv2.resize(image, (w, h), interpolation=cv2.INTER_LANCZOS4)
+    return image
+
+
+# make uc and prompt shapes match via padding for long prompts
+null_cond = None
+
+def fix_cond_shapes(model, prompt_condition, uc):
+    if uc is None:
+        return prompt_condition, uc
+    global null_cond
+    if null_cond is None:
+        null_cond = model.get_learned_conditioning([""])
+    while prompt_condition.shape[1] > uc.shape[1]:
+        uc = torch.cat((uc, null_cond.repeat((uc.shape[0], 1, 1))), axis=1)
+    while prompt_condition.shape[1] < uc.shape[1]:
+        prompt_condition = torch.cat((prompt_condition, null_cond.repeat((prompt_condition.shape[0], 1, 1))), axis=1)
+    return prompt_condition, uc

libcom/painterly_image_harmonization/source/PHDiffusion/stable_diffusion.yaml

@@ -0,0 +1,75 @@
+name: test
+model:
+  base_learning_rate: 1.0e-04
+  target: libcom.painterly_image_harmonization.source.PHDiffusion.ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: False
+
+    scheduler_config: # 10000 warmup steps
+      target: libcom.painterly_image_harmonization.source.PHDiffusion.ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 10000 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: libcom.painterly_image_harmonization.source.PHDiffusion.ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: libcom.painterly_image_harmonization.source.PHDiffusion.ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 512
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config: #__is_unconditional__
+      target: libcom.painterly_image_harmonization.source.PHDiffusion.ldm.modules.encoders.modules.FrozenCLIPEmbedder
+      params:
+        model_path: openai-clip-vit-large-patch14
+
+

libcom/reflection_generation/__init__.py

@@ -0,0 +1,3 @@
+from .reflection_generation import ReflectionGenerationModel
+
+__all__ = ['ReflectionGenerationModel']

libcom/reflection_generation/reflection_generation.py

@@ -0,0 +1,174 @@
+import torch
+import torchvision.transforms as transforms
+from libcom.utils.model_download import download_pretrained_model, download_entire_folder
+from libcom.utils.process_image import *
+from libcom.utils.environment import *
+import os
+from omegaconf import OmegaConf
+from pytorch_lightning import seed_everything
+import numpy as np
+import os
+import cv2
+from PIL import Image
+import numpy as np
+
+try:
+    from lightning_fabric.utilities.seed import log
+    log.propagate = False
+except:
+    pass
+from .source.PostProcessModel import PostProcess
+from .source.cldm.model import load_state_dict
+
+cur_dir   = os.path.dirname(os.path.abspath(__file__))
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+model_set = ['ReflectionGenerationModel']
+
+class ReflectionGenerationModel:
+    """
+    Foreground reflection generation model based on diffusion model and control net.
+
+    Args:
+        device (str | torch.device): gpu id
+        model_type (str): predefined model type
+        kwargs (dict): other parameters for building model
+
+    Examples:
+        >>> from libcom import ReflectionGenerationModel
+        >>> from libcom.utils.process_image import make_image_grid
+        >>> import cv2
+        >>> net = ReflectionGenerationModel(device=2, model_type='ReflectionGeneration')
+        >>> comp_image1 = "../tests/reflection_generation/composite/1.png"
+        >>> comp_mask1 = "../tests/reflection_generation/composite_mask/1.png"
+        >>> preds = net(comp_image1, comp_mask1, number=5)
+        >>> grid_img  = make_image_grid([comp_image1, comp_mask1] + preds)
+        >>> cv2.imwrite('../docs/_static/image/reflection_generation_result1.jpg', grid_img)
+        >>> comp_image2 = "../tests/reflection_generation/composite/2.png"
+        >>> comp_mask2 = "../tests/reflection_generation/composite_mask/2.png"
+        >>> preds = net(comp_image2, comp_mask2, number=5)
+        >>> grid_img  = make_image_grid([comp_image2, comp_mask2] + preds)
+        >>> cv2.imwrite('../docs/_static/image/reflection_generation_result2.jpg', grid_img)
+
+    Expected result:
+
+    .. image:: _static/image/reflection_generation1.jpg
+        :scale: 21 %
+    .. image:: _static/image/reflection_generation2.jpg
+        :scale: 21 %
+
+        
+    """
+    def __init__(self, device=0, model_type='ReflectionGeneration', **kwargs):
+        # assert model_type in model_set, f'Not implementation for {model_type}'
+        self.model_type = model_type
+        self.option = kwargs
+        
+        cldm_weight_path = os.path.join(model_dir, 'pretrained_models', 'Reflection_cldm.ckpt')
+        ppp_weight_path = os.path.join(model_dir, 'pretrained_models', 'Reflection_ppp.ckpt')
+        reg_net_path = os.path.join(model_dir, 'pretrained_models', 'Reflection_reg.pth')
+        
+        download_pretrained_model(cldm_weight_path)
+        download_pretrained_model(ppp_weight_path)
+        download_pretrained_model(reg_net_path)
+        
+        self.device = check_gpu_device(device)
+        self.build_pretrained_model(ppp_weight_path, cldm_weight_path, reg_net_path)
+        self.build_data_transformer()
+
+    def build_pretrained_model(self, ppp_weight_path, cldm_weight_path, reg_weight_path):
+        config_path = '../libcom/reflection_generation/source/cldm_v15.yaml'
+        config = OmegaConf.load(config_path)
+        config.model.params.reg_net_path = reg_weight_path
+        clip_path = os.path.join(model_dir, '../shared_pretrained_models', 'openai-clip-vit-large-patch14')
+        download_entire_folder(clip_path)
+        config.model.params.cond_stage_config.params.version = clip_path    
+        model = PostProcess(
+            model_path=config,
+            control_net_path=cldm_weight_path,
+            infe_steps=50
+            )
+        model.load_state_dict(load_state_dict(ppp_weight_path, location='cpu'), strict=False)
+        self.model = model.to(self.device).eval()
+        
+    def build_data_transformer(self):
+        self.image_size = 512
+        self.transformer = transforms.Compose([
+            transforms.Resize((self.image_size, self.image_size)),
+            transforms.ToTensor(),
+        ])
+    
+    def inputs_preprocess(self, composite_image, composite_mask):
+        img = read_image_pil(composite_image)
+        img = self.transformer(img).permute(1, 2, 0)
+        target = img * 2 - 1
+
+        mask = read_mask_pil(composite_mask)
+        mask_np = np.array(mask)
+        mask_np = cv2.resize(mask_np, (512, 512))
+        _, fg_instance_thresh = cv2.threshold(mask_np, 128, 255, cv2.THRESH_BINARY)
+        contours_instance, _ = cv2.findContours(fg_instance_thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
+        merged_contour_points_instance = np.concatenate(contours_instance)
+        rect_instance = cv2.minAreaRect(merged_contour_points_instance)
+        (x, y), (w, h), theta = rect_instance
+        if w < h:
+            w, h = h, w
+            theta += 90
+        bbx_instance = np.array([x, y, w + 1, h + 1, theta]).astype(int)
+        bbx_instance = torch.tensor(bbx_instance).unsqueeze(0)
+
+        mask = self.transformer(mask).permute(1, 2, 0)
+        img, mask, target = img.unsqueeze(0), mask.unsqueeze(0), target.unsqueeze(0)
+
+        cat_img = torch.cat([img, mask], dim=-1)
+        mask_embeddings = torch.zeros((1, 64, 2048), dtype=torch.float32)
+        bbx_region = torch.zeros((1, 512, 512), dtype=torch.float32)
+
+        object_mask = mask[:,:,:,0]
+
+        return target.to(self.device), mask.to(self.device), cat_img.to(self.device), object_mask.to(self.device), \
+               bbx_instance.to(self.device), mask_embeddings.to(self.device), bbx_region.to(self.device)
+    
+    def outputs_postprocess(self, outputs):
+        output = outputs
+        adjusted_img = output[:, :, :, :3]
+        adjusted_img = torch.clamp(adjusted_img, -1., 1.)
+        adjusted_img = (adjusted_img + 1.0) / 2.0
+        adjusted_img = (adjusted_img * 255).int()
+        composite_img = adjusted_img
+        composite_img = np.array(composite_img.cpu().squeeze(0), dtype=np.uint8)
+        composite_img = composite_img[:,:,[2,1,0]]
+        return composite_img
+    
+    @torch.no_grad()
+    def inf_img(self, inputs):
+        target, mask, cat_img, object_mask, bbx_instance, mask_embeddings, bbx_region = inputs
+        batch = dict(jpg=target, cls=cat_img, objectmask=object_mask, fg=bbx_instance, bbx=bbx_region, embeddings=mask_embeddings, txt=[''], hint=cat_img)
+        images = self.model.model.log_images(batch, use_x_T=True)
+        output = images['samples_cfg_scale_9.00'].permute(0,2,3,1)
+
+        return output
+    
+
+    @torch.no_grad()
+    def __call__(self, composite_image, composite_mask, number=5, seed=42):
+        """
+        Generate reflection for foreground object.
+        
+        Args:
+            composite_img (str | numpy.ndarray): The path to composite image or composite image in ndarray form.
+            composite_mask (str | numpy.ndarray): The path to foreground object mask or foreground object mask in ndarray form.
+            number (int): Number of images to be inferenced. default: 5.
+            seed: Random Seed is used to reproduce results and same seed will lead to same results. 
+        
+        Returns: 
+            generated_images (list): A list of images with generated foreground reflections. Each image is in ndarray form with a shape of 512x512x3
+
+        """
+        seed_everything(seed)
+        inputs = self.inputs_preprocess(composite_image, composite_mask)
+        preds = []
+        for _ in range(number):
+            outputs = self.inf_img(inputs)
+            pred = self.outputs_postprocess(outputs)
+            preds.append(pred)
+        return preds