libcom 0.1.8__py3-none-any.whl

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/modules.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+import math
+from torch.utils.checkpoint import checkpoint
+
+from transformers import CLIPTokenizer, CLIPModel
+
+import open_clip
+import re
+from libcom.painterly_image_harmonization.source.PHDiffusion.ldm.util import default, count_params
+
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class IdentityEncoder(AbstractEncoder):
+
+    def encode(self, x):
+        return x
+
+
+class ClassEmbedder(nn.Module):
+    def __init__(self, embed_dim, n_classes=1000, key='class'):
+        super().__init__()
+        self.key = key
+        self.embedding = nn.Embedding(n_classes, embed_dim)
+
+    def forward(self, batch, key=None):
+        if key is None:
+            key = self.key
+        # this is for use in crossattn
+        c = batch[key][:, None]
+        c = self.embedding(c)
+        return c
+
+
+class FrozenCLIPEmbedder(AbstractEncoder):
+    """Uses the CLIP transformer encoder for text (from huggingface)"""
+    def __init__(self, model_path="openai/clip-vit-large-patch14", device="cuda", max_length=77,
+                 freeze=True, layer="last"):  # clip-vit-base-patch32
+        super().__init__()
+        import os
+        self.tokenizer   = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path=model_path)
+        self.transformer = CLIPModel.from_pretrained(pretrained_model_name_or_path=model_path).text_model
+        self.device      = device
+        self.max_length  = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        device = next(self.transformer.parameters()).device
+        tokens = batch_encoding["input_ids"].to(device)
+        outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer != 'last')
+
+        if self.layer == 'penultimate':
+            z = outputs.hidden_states[-2]
+            z = self.transformer.final_layer_norm(z)
+        else:
+            z = outputs.last_hidden_state
+        return z
+
+    def encode(self, text):
+        return self(text)
+
+
+
+if __name__ == "__main__":
+    model = FrozenCLIPEmbedder()
+    count_params(model, verbose=True)

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/api.py

@@ -0,0 +1,269 @@
+from enum import Enum, unique
+
+import cv2
+import torch
+from basicsr.utils import img2tensor
+from libcom.painterly_image_harmonization.source.ldm.util import resize_numpy_image
+from PIL import Image
+from torch import autocast
+
+
+@unique
+class ExtraCondition(Enum):
+    sketch = 0
+    keypose = 1
+    seg = 2
+    depth = 3
+    canny = 4
+    style = 5
+    color = 6
+    openpose = 7
+
+
+def get_cond_model(opt, cond_type: ExtraCondition):
+    if cond_type == ExtraCondition.sketch:
+        from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.model_edge import pidinet
+        model = pidinet()
+        ckp = torch.load('models/table5_pidinet.pth', map_location='cpu')['state_dict']
+        model.load_state_dict({k.replace('module.', ''): v for k, v in ckp.items()}, strict=True)
+        model.to(opt.device)
+        return model
+    elif cond_type == ExtraCondition.seg:
+        raise NotImplementedError
+    elif cond_type == ExtraCondition.keypose:
+        import mmcv
+        from mmdet.apis import init_detector
+        from mmpose.apis import init_pose_model
+        det_config = 'configs/mm/faster_rcnn_r50_fpn_coco.py'
+        det_checkpoint = 'models/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'
+        pose_config = 'configs/mm/hrnet_w48_coco_256x192.py'
+        pose_checkpoint = 'models/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth'
+        det_config_mmcv = mmcv.Config.fromfile(det_config)
+        det_model = init_detector(det_config_mmcv, det_checkpoint, device=opt.device)
+        pose_config_mmcv = mmcv.Config.fromfile(pose_config)
+        pose_model = init_pose_model(pose_config_mmcv, pose_checkpoint, device=opt.device)
+        return {'pose_model': pose_model, 'det_model': det_model}
+    elif cond_type == ExtraCondition.depth:
+        from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.midas.api import MiDaSInference
+        model = MiDaSInference(model_type='dpt_hybrid').to(opt.device)
+        return model
+    elif cond_type == ExtraCondition.canny:
+        return None
+    elif cond_type == ExtraCondition.style:
+        from transformers import CLIPProcessor, CLIPVisionModel
+        version = 'openai/clip-vit-large-patch14'
+        processor = CLIPProcessor.from_pretrained(version)
+        clip_vision_model = CLIPVisionModel.from_pretrained(version).to(opt.device)
+        return {'processor': processor, 'clip_vision_model': clip_vision_model}
+    elif cond_type == ExtraCondition.color:
+        return None
+    elif cond_type == ExtraCondition.openpose:
+        from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.openpose.api import OpenposeInference
+        model = OpenposeInference().to(opt.device)
+        return model
+    else:
+        raise NotImplementedError
+
+
+def get_cond_sketch(opt, cond_image, cond_inp_type, cond_model=None):
+    if isinstance(cond_image, str):
+        edge = cv2.imread(cond_image)
+    else:
+        # for gradio input, pay attention, it's rgb numpy
+        edge = cv2.cvtColor(cond_image, cv2.COLOR_RGB2BGR)
+    edge = resize_numpy_image(edge, max_resolution=opt.max_resolution, resize_short_edge=opt.resize_short_edge)
+    opt.H, opt.W = edge.shape[:2]
+    if cond_inp_type == 'sketch':
+        edge = img2tensor(edge)[0].unsqueeze(0).unsqueeze(0) / 255.
+        edge = edge.to(opt.device)
+    elif cond_inp_type == 'image':
+        edge = img2tensor(edge).unsqueeze(0) / 255.
+        edge = cond_model(edge.to(opt.device))[-1]
+    else:
+        raise NotImplementedError
+
+    # edge = 1-edge # for white background
+    edge = edge > 0.5
+    edge = edge.float()
+
+    return edge
+
+
+def get_cond_seg(opt, cond_image, cond_inp_type='image', cond_model=None):
+    if isinstance(cond_image, str):
+        seg = cv2.imread(cond_image)
+    else:
+        seg = cv2.cvtColor(cond_image, cv2.COLOR_RGB2BGR)
+    seg = resize_numpy_image(seg, max_resolution=opt.max_resolution, resize_short_edge=opt.resize_short_edge)
+    opt.H, opt.W = seg.shape[:2]
+    if cond_inp_type == 'seg':
+        seg = img2tensor(seg).unsqueeze(0) / 255.
+        seg = seg.to(opt.device)
+    else:
+        raise NotImplementedError
+
+    return seg
+
+
+def get_cond_keypose(opt, cond_image, cond_inp_type='image', cond_model=None):
+    if isinstance(cond_image, str):
+        pose = cv2.imread(cond_image)
+    else:
+        pose = cv2.cvtColor(cond_image, cv2.COLOR_RGB2BGR)
+    pose = resize_numpy_image(pose, max_resolution=opt.max_resolution, resize_short_edge=opt.resize_short_edge)
+    opt.H, opt.W = pose.shape[:2]
+    if cond_inp_type == 'keypose':
+        pose = img2tensor(pose).unsqueeze(0) / 255.
+        pose = pose.to(opt.device)
+    elif cond_inp_type == 'image':
+        from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.utils import imshow_keypoints
+        from mmdet.apis import inference_detector
+        from mmpose.apis import (inference_top_down_pose_model, process_mmdet_results)
+
+        # mmpose seems not compatible with autocast fp16
+        with autocast("cuda", dtype=torch.float32):
+            mmdet_results = inference_detector(cond_model['det_model'], pose)
+            # keep the person class bounding boxes.
+            person_results = process_mmdet_results(mmdet_results, 1)
+
+            # optional
+            return_heatmap = False
+            dataset = cond_model['pose_model'].cfg.data['test']['type']
+
+            # e.g. use ('backbone', ) to return backbone feature
+            output_layer_names = None
+            pose_results, returned_outputs = inference_top_down_pose_model(
+                cond_model['pose_model'],
+                pose,
+                person_results,
+                bbox_thr=0.2,
+                format='xyxy',
+                dataset=dataset,
+                dataset_info=None,
+                return_heatmap=return_heatmap,
+                outputs=output_layer_names)
+
+        # show the results
+        pose = imshow_keypoints(pose, pose_results, radius=2, thickness=2)
+        pose = img2tensor(pose).unsqueeze(0) / 255.
+        pose = pose.to(opt.device)
+    else:
+        raise NotImplementedError
+
+    return pose
+
+
+def get_cond_depth(opt, cond_image, cond_inp_type='image', cond_model=None):
+    if isinstance(cond_image, str):
+        depth = cv2.imread(cond_image)
+    else:
+        depth = cv2.cvtColor(cond_image, cv2.COLOR_RGB2BGR)
+    depth = resize_numpy_image(depth, max_resolution=opt.max_resolution, resize_short_edge=opt.resize_short_edge)
+    opt.H, opt.W = depth.shape[:2]
+    if cond_inp_type == 'depth':
+        depth = img2tensor(depth).unsqueeze(0) / 255.
+        depth = depth.to(opt.device)
+    elif cond_inp_type == 'image':
+        depth = img2tensor(depth).unsqueeze(0) / 127.5 - 1.0
+        depth = cond_model(depth.to(opt.device)).repeat(1, 3, 1, 1)
+        depth -= torch.min(depth)
+        depth /= torch.max(depth)
+    else:
+        raise NotImplementedError
+
+    return depth
+
+
+def get_cond_canny(opt, cond_image, cond_inp_type='image', cond_model=None):
+    if isinstance(cond_image, str):
+        canny = cv2.imread(cond_image)
+    else:
+        canny = cv2.cvtColor(cond_image, cv2.COLOR_RGB2BGR)
+    canny = resize_numpy_image(canny, max_resolution=opt.max_resolution, resize_short_edge=opt.resize_short_edge)
+    opt.H, opt.W = canny.shape[:2]
+    if cond_inp_type == 'canny':
+        canny = img2tensor(canny)[0:1].unsqueeze(0) / 255.
+        canny = canny.to(opt.device)
+    elif cond_inp_type == 'image':
+        canny = cv2.Canny(canny, 100, 200)[..., None]
+        canny = img2tensor(canny).unsqueeze(0) / 255.
+        canny = canny.to(opt.device)
+    else:
+        raise NotImplementedError
+
+    return canny
+
+
+def get_cond_style(opt, cond_image, cond_inp_type='image', cond_model=None):
+    assert cond_inp_type == 'image'
+    if isinstance(cond_image, str):
+        style = Image.open(cond_image)
+    else:
+        # numpy image to PIL image
+        style = Image.fromarray(cond_image)
+
+    style_for_clip = cond_model['processor'](images=style, return_tensors="pt")['pixel_values']
+    style_feat = cond_model['clip_vision_model'](style_for_clip.to(opt.device))['last_hidden_state']
+
+    return style_feat
+
+
+def get_cond_color(opt, cond_image, cond_inp_type='image', cond_model=None):
+    if isinstance(cond_image, str):
+        color = cv2.imread(cond_image)
+    else:
+        color = cv2.cvtColor(cond_image, cv2.COLOR_RGB2BGR)
+    color = resize_numpy_image(color, max_resolution=opt.max_resolution, resize_short_edge=opt.resize_short_edge)
+    opt.H, opt.W = color.shape[:2]
+    if cond_inp_type == 'image':
+        color = cv2.resize(color, (opt.W//64, opt.H//64), interpolation=cv2.INTER_CUBIC)
+        color = cv2.resize(color, (opt.W, opt.H), interpolation=cv2.INTER_NEAREST)
+    color = img2tensor(color).unsqueeze(0) / 255.
+    color = color.to(opt.device)
+    return color
+
+
+def get_cond_openpose(opt, cond_image, cond_inp_type='image', cond_model=None):
+    if isinstance(cond_image, str):
+        openpose_keypose = cv2.imread(cond_image)
+    else:
+        openpose_keypose = cv2.cvtColor(cond_image, cv2.COLOR_RGB2BGR)
+    openpose_keypose = resize_numpy_image(
+        openpose_keypose, max_resolution=opt.max_resolution, resize_short_edge=opt.resize_short_edge)
+    opt.H, opt.W = openpose_keypose.shape[:2]
+    if cond_inp_type == 'openpose':
+        openpose_keypose = img2tensor(openpose_keypose).unsqueeze(0) / 255.
+        openpose_keypose = openpose_keypose.to(opt.device)
+    elif cond_inp_type == 'image':
+        with autocast('cuda', dtype=torch.float32):
+            openpose_keypose = cond_model(openpose_keypose)
+        openpose_keypose = img2tensor(openpose_keypose).unsqueeze(0) / 255.
+        openpose_keypose = openpose_keypose.to(opt.device)
+
+    else:
+        raise NotImplementedError
+
+    return openpose_keypose
+
+
+def get_adapter_feature(inputs, adapters):
+    ret_feat_map = None
+    ret_feat_seq = None
+    if not isinstance(inputs, list):
+        inputs = [inputs]
+        adapters = [adapters]
+
+    for input, adapter in zip(inputs, adapters):
+        cur_feature = adapter['model'](input)
+        if isinstance(cur_feature, list):
+            if ret_feat_map is None:
+                ret_feat_map = list(map(lambda x: x * adapter['cond_weight'], cur_feature))
+            else:
+                ret_feat_map = list(map(lambda x, y: x + y * adapter['cond_weight'], ret_feat_map, cur_feature))
+        else:
+            if ret_feat_seq is None:
+                ret_feat_seq = cur_feature * adapter['cond_weight']
+            else:
+                ret_feat_seq = torch.cat([ret_feat_seq, cur_feature * adapter['cond_weight']], dim=1)
+
+    return ret_feat_map, ret_feat_seq

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/api.py

@@ -0,0 +1,175 @@
+# based on https://github.com/isl-org/MiDaS
+import os
+
+import cv2
+import torch
+import torch.nn as nn
+from torchvision.transforms import Compose
+
+from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.midas.midas.dpt_depth import DPTDepthModel
+from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.midas.midas.midas_net import MidasNet
+from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.midas.midas.midas_net_custom import MidasNet_small
+from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.midas.midas.transforms import Resize, NormalizeImage, PrepareForNet
+
+
+ISL_PATHS = {
+    "dpt_large": "models/dpt_large-midas-2f21e586.pt",
+    "dpt_hybrid": "models/dpt_hybrid-midas-501f0c75.pt",
+    "midas_v21": "",
+    "midas_v21_small": "",
+}
+
+remote_model_path = "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt"
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+def load_midas_transform(model_type):
+    # https://github.com/isl-org/MiDaS/blob/master/run.py
+    # load transform only
+    if model_type == "dpt_large":  # DPT-Large
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "dpt_hybrid":  # DPT-Hybrid
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "midas_v21":
+        net_w, net_h = 384, 384
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    elif model_type == "midas_v21_small":
+        net_w, net_h = 256, 256
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    else:
+        assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
+
+    transform = Compose(
+        [
+            Resize(
+                net_w,
+                net_h,
+                resize_target=None,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=32,
+                resize_method=resize_mode,
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            normalization,
+            PrepareForNet(),
+        ]
+    )
+
+    return transform
+
+
+def load_model(model_type):
+    # https://github.com/isl-org/MiDaS/blob/master/run.py
+    # load network
+    model_path = ISL_PATHS[model_type]
+    if model_type == "dpt_large":  # DPT-Large
+        model = DPTDepthModel(
+            path=model_path,
+            backbone="vitl16_384",
+            non_negative=True,
+        )
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "dpt_hybrid":  # DPT-Hybrid
+        if not os.path.exists(model_path):
+            from basicsr.utils.download_util import load_file_from_url
+            load_file_from_url(remote_model_path, model_dir='models')
+
+        model = DPTDepthModel(
+            path=model_path,
+            backbone="vitb_rn50_384",
+            non_negative=True,
+        )
+        net_w, net_h = 384, 384
+        resize_mode = "minimal"
+        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+
+    elif model_type == "midas_v21":
+        model = MidasNet(model_path, non_negative=True)
+        net_w, net_h = 384, 384
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+
+    elif model_type == "midas_v21_small":
+        model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True,
+                               non_negative=True, blocks={'expand': True})
+        net_w, net_h = 256, 256
+        resize_mode = "upper_bound"
+        normalization = NormalizeImage(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+
+    else:
+        print(f"model_type '{model_type}' not implemented, use: --model_type large")
+        assert False
+
+    transform = Compose(
+        [
+            Resize(
+                net_w,
+                net_h,
+                resize_target=None,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=32,
+                resize_method=resize_mode,
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            normalization,
+            PrepareForNet(),
+        ]
+    )
+
+    return model.eval(), transform
+
+
+class MiDaSInference(nn.Module):
+    MODEL_TYPES_TORCH_HUB = [
+        "DPT_Large",
+        "DPT_Hybrid",
+        "MiDaS_small"
+    ]
+    MODEL_TYPES_ISL = [
+        "dpt_large",
+        "dpt_hybrid",
+        "midas_v21",
+        "midas_v21_small",
+    ]
+
+    def __init__(self, model_type):
+        super().__init__()
+        assert (model_type in self.MODEL_TYPES_ISL)
+        model, _ = load_model(model_type)
+        self.model = model
+        self.model.train = disabled_train
+
+    def forward(self, x):
+        # x in 0..1 as produced by calling self.transform on a 0..1 float64 numpy array
+        # NOTE: we expect that the correct transform has been called during dataloading.
+        with torch.no_grad():
+            prediction = self.model(x)
+            prediction = torch.nn.functional.interpolate(
+                prediction.unsqueeze(1),
+                size=x.shape[2:],
+                mode="bicubic",
+                align_corners=False,
+            )
+        assert prediction.shape == (x.shape[0], 1, x.shape[2], x.shape[3])
+        return prediction

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/base_model.py

@@ -0,0 +1,16 @@
+import torch
+
+
+class BaseModel(torch.nn.Module):
+    def load(self, path):
+        """Load model from file.
+
+        Args:
+            path (str): file path
+        """
+        parameters = torch.load(path, map_location=torch.device('cpu'))
+
+        if "optimizer" in parameters:
+            parameters = parameters["model"]
+
+        self.load_state_dict(parameters)