Upload folder using huggingface_hub

DensePose/__init__.py

@@ -0,0 +1,158 @@
+
+import glob
+import os
+from random import randint
+import shutil
+import time
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+from densepose import add_densepose_config
+from densepose.vis.base import CompoundVisualizer
+from densepose.vis.densepose_results import DensePoseResultsFineSegmentationVisualizer
+from densepose.vis.extractor import create_extractor, CompoundExtractor
+from detectron2.config import get_cfg
+from detectron2.data.detection_utils import read_image
+from detectron2.engine.defaults import DefaultPredictor
+
+
+class DensePose:
+    """
+    DensePose used in this project is from Detectron2 (https://github.com/facebookresearch/detectron2).
+    These codes are modified from https://github.com/facebookresearch/detectron2/tree/main/projects/DensePose.
+    The checkpoint is downloaded from https://github.com/facebookresearch/detectron2/blob/main/projects/DensePose/doc/DENSEPOSE_IUV.md#ModelZoo.
+
+    We use the model R_50_FPN_s1x with id 165712039, but other models should also work.
+    The config file is downloaded from https://github.com/facebookresearch/detectron2/tree/main/projects/DensePose/configs.
+    Noted that the config file should match the model checkpoint and Base-DensePose-RCNN-FPN.yaml is also needed.
+    """
+
+    def __init__(self, model_path="./checkpoints/densepose_", device="cuda"):
+        self.device = device
+        self.config_path = os.path.join(model_path, 'densepose_rcnn_R_50_FPN_s1x.yaml')
+        self.model_path = os.path.join(model_path, 'model_final_162be9.pkl')
+        self.visualizations = ["dp_segm"]
+        self.VISUALIZERS = {"dp_segm": DensePoseResultsFineSegmentationVisualizer}
+        self.min_score = 0.8
+
+        self.cfg = self.setup_config()
+        self.predictor = DefaultPredictor(self.cfg)
+        self.predictor.model.to(self.device)
+
+    def setup_config(self):
+        opts = ["MODEL.ROI_HEADS.SCORE_THRESH_TEST", str(self.min_score)]
+        cfg = get_cfg()
+        add_densepose_config(cfg)
+        cfg.merge_from_file(self.config_path)
+        cfg.merge_from_list(opts)
+        cfg.MODEL.WEIGHTS = self.model_path
+        cfg.freeze()
+        return cfg
+
+    @staticmethod
+    def _get_input_file_list(input_spec: str):
+        if os.path.isdir(input_spec):
+            file_list = [os.path.join(input_spec, fname) for fname in os.listdir(input_spec)
+                         if os.path.isfile(os.path.join(input_spec, fname))]
+        elif os.path.isfile(input_spec):
+            file_list = [input_spec]
+        else:
+            file_list = glob.glob(input_spec)
+        return file_list
+
+    def create_context(self, cfg, output_path):
+        vis_specs = self.visualizations
+        visualizers = []
+        extractors = []
+        for vis_spec in vis_specs:
+            texture_atlas = texture_atlases_dict = None
+            vis = self.VISUALIZERS[vis_spec](
+                cfg=cfg,
+                texture_atlas=texture_atlas,
+                texture_atlases_dict=texture_atlases_dict,
+                alpha=1.0
+            )
+            visualizers.append(vis)
+            extractor = create_extractor(vis)
+            extractors.append(extractor)
+        visualizer = CompoundVisualizer(visualizers)
+        extractor = CompoundExtractor(extractors)
+        context = {
+            "extractor": extractor,
+            "visualizer": visualizer,
+            "out_fname": output_path,
+            "entry_idx": 0,
+        }
+        return context
+
+    def execute_on_outputs(self, context, entry, outputs):
+        extractor = context["extractor"]
+
+        data = extractor(outputs)
+
+        H, W, _ = entry["image"].shape
+        result = np.zeros((H, W), dtype=np.uint8)
+
+        data, box = data[0]
+        x, y, w, h = [int(_) for _ in box[0].cpu().numpy()]
+        i_array = data[0].labels[None].cpu().numpy()[0]
+        result[y:y + h, x:x + w] = i_array
+        result = Image.fromarray(result)
+        result.save(context["out_fname"])
+
+    def __call__(self, image_or_path, resize=512) -> Image.Image:
+        """
+        :param image_or_path: Path of the input image.
+        :param resize: Resize the input image if its max size is larger than this value.
+        :return: Dense pose image.
+        """
+        # random tmp path with timestamp
+        tmp_path = f"./densepose_/tmp/"
+        if not os.path.exists(tmp_path):
+            os.makedirs(tmp_path)
+
+        image_path = os.path.join(tmp_path, f"{int(time.time())}-{self.device}-{randint(0, 100000)}.png")
+        if isinstance(image_or_path, str):
+            assert image_or_path.split(".")[-1] in ["jpg", "png"], "Only support jpg and png images."
+            shutil.copy(image_or_path, image_path)
+        elif isinstance(image_or_path, Image.Image):
+            image_or_path.save(image_path)
+        else:
+            shutil.rmtree(tmp_path)
+            raise TypeError("image_path must be str or PIL.Image.Image")
+
+        output_path = image_path.replace(".png", "_dense.png").replace(".jpg", "_dense.png")
+        w, h = Image.open(image_path).size
+
+        file_list = self._get_input_file_list(image_path)
+        assert len(file_list), "No input images found!"
+        context = self.create_context(self.cfg, output_path)
+        for file_name in file_list:
+            img = read_image(file_name, format="BGR")  # predictor expects BGR image.
+            # resize
+            if (_ := max(img.shape)) > resize:
+                scale = resize / _
+                img = cv2.resize(img, (int(img.shape[1] * scale), int(img.shape[0] * scale)))
+
+            with torch.no_grad():
+                outputs = self.predictor(img)["instances"]
+                try:
+                    self.execute_on_outputs(context, {"file_name": file_name, "image": img}, outputs)
+                except Exception as e:
+                    null_gray = Image.new('L', (1, 1))
+                    null_gray.save(output_path)
+
+        dense_gray = Image.open(output_path).convert("L")
+        dense_gray = dense_gray.resize((w, h), Image.NEAREST)
+        # remove image_path and output_path
+        os.remove(image_path)
+        os.remove(output_path)
+
+
+        return dense_gray
+
+
+if __name__ == '__main__':
+    pass

SCHP/__init__.py

@@ -0,0 +1,179 @@
+from model.SCHP import networks
+from model.SCHP.utils.transforms import get_affine_transform, transform_logits
+
+from collections import OrderedDict
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from torchvision import transforms
+
+def get_palette(num_cls):
+    """ Returns the color map for visualizing the segmentation mask.
+    Args:
+        num_cls: Number of classes
+    Returns:
+        The color map
+    """
+    n = num_cls
+    palette = [0] * (n * 3)
+    for j in range(0, n):
+        lab = j
+        palette[j * 3 + 0] = 0
+        palette[j * 3 + 1] = 0
+        palette[j * 3 + 2] = 0
+        i = 0
+        while lab:
+            palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
+            palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
+            palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
+            i += 1
+            lab >>= 3
+    return palette
+
+dataset_settings = {
+    'lip': {
+        'input_size': [473, 473],
+        'num_classes': 20,
+        'label': ['Background', 'Hat', 'Hair', 'Glove', 'Sunglasses', 'Upper-clothes', 'Dress', 'Coat',
+                  'Socks', 'Pants', 'Jumpsuits', 'Scarf', 'Skirt', 'Face', 'Left-arm', 'Right-arm',
+                  'Left-leg', 'Right-leg', 'Left-shoe', 'Right-shoe']
+    },
+    'atr': {
+        'input_size': [512, 512],
+        'num_classes': 18,
+        'label': ['Background', 'Hat', 'Hair', 'Sunglasses', 'Upper-clothes', 'Skirt', 'Pants', 'Dress', 'Belt',
+                  'Left-shoe', 'Right-shoe', 'Face', 'Left-leg', 'Right-leg', 'Left-arm', 'Right-arm', 'Bag', 'Scarf']
+    },
+    'pascal': {
+        'input_size': [512, 512],
+        'num_classes': 7,
+        'label': ['Background', 'Head', 'Torso', 'Upper Arms', 'Lower Arms', 'Upper Legs', 'Lower Legs'],
+    }
+}
+
+class SCHP:
+    def __init__(self, ckpt_path, device):
+        dataset_type = None
+        if 'lip' in ckpt_path:
+            dataset_type = 'lip'
+        elif 'atr' in ckpt_path:
+            dataset_type = 'atr'
+        elif 'pascal' in ckpt_path:
+            dataset_type = 'pascal'
+        assert dataset_type is not None, 'Dataset type not found in checkpoint path'
+        self.device = device
+        self.num_classes = dataset_settings[dataset_type]['num_classes']
+        self.input_size = dataset_settings[dataset_type]['input_size']
+        self.aspect_ratio = self.input_size[1] * 1.0 / self.input_size[0]
+        self.palette = get_palette(self.num_classes)
+
+        self.label = dataset_settings[dataset_type]['label']
+        self.model = networks.init_model('resnet101', num_classes=self.num_classes, pretrained=None).to(device)
+        self.load_ckpt(ckpt_path)
+        self.model.eval()
+        
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229])
+        ])
+        self.upsample = torch.nn.Upsample(size=self.input_size, mode='bilinear', align_corners=True)
+
+
+    def load_ckpt(self, ckpt_path):
+        rename_map = {
+            "decoder.conv3.2.weight": "decoder.conv3.3.weight",
+            "decoder.conv3.3.weight": "decoder.conv3.4.weight",
+            "decoder.conv3.3.bias": "decoder.conv3.4.bias",
+            "decoder.conv3.3.running_mean": "decoder.conv3.4.running_mean",
+            "decoder.conv3.3.running_var": "decoder.conv3.4.running_var",
+            "fushion.3.weight": "fushion.4.weight",
+            "fushion.3.bias": "fushion.4.bias",
+        }
+        state_dict = torch.load(ckpt_path, map_location='cpu')['state_dict']
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            name = k[7:]  # remove `module.`
+            new_state_dict[name] = v
+        new_state_dict_ = OrderedDict()
+        for k, v in list(new_state_dict.items()):
+            if k in rename_map:
+                new_state_dict_[rename_map[k]] = v
+            else:
+                new_state_dict_[k] = v
+        self.model.load_state_dict(new_state_dict_, strict=False)
+
+    def _box2cs(self, box):
+        x, y, w, h = box[:4]
+        return self._xywh2cs(x, y, w, h)
+
+    def _xywh2cs(self, x, y, w, h):
+        center = np.zeros((2), dtype=np.float32)
+        center[0] = x + w * 0.5
+        center[1] = y + h * 0.5
+        if w > self.aspect_ratio * h:
+            h = w * 1.0 / self.aspect_ratio
+        elif w < self.aspect_ratio * h:
+            w = h * self.aspect_ratio
+        scale = np.array([w, h], dtype=np.float32)
+        return center, scale
+
+    def preprocess(self, image):
+        if isinstance(image, str):
+            img = cv2.imread(image, cv2.IMREAD_COLOR)
+        elif isinstance(image, Image.Image):
+            # to cv2 format
+            img = np.array(image)
+    
+        h, w, _ = img.shape
+        # Get person center and scale
+        person_center, s = self._box2cs([0, 0, w - 1, h - 1])
+        r = 0
+        trans = get_affine_transform(person_center, s, r, self.input_size)
+        input = cv2.warpAffine(
+            img,
+            trans,
+            (int(self.input_size[1]), int(self.input_size[0])),
+            flags=cv2.INTER_LINEAR,
+            borderMode=cv2.BORDER_CONSTANT,
+            borderValue=(0, 0, 0))
+
+        input = self.transform(input).to(self.device).unsqueeze(0)
+        meta = {
+                'center': person_center,
+                'height': h,
+                'width': w,
+                'scale': s,
+                'rotation': r
+        }
+        return input, meta
+
+
+    def __call__(self, image_or_path):
+        if isinstance(image_or_path, list):
+            image_list = []
+            meta_list = []
+            for image in image_or_path:
+                image, meta = self.preprocess(image)
+                image_list.append(image)
+                meta_list.append(meta)
+            image = torch.cat(image_list, dim=0)
+        else:
+            image, meta = self.preprocess(image_or_path)
+            meta_list = [meta]
+                
+        output = self.model(image)
+        # upsample_outputs = self.upsample(output[0][-1])
+        upsample_outputs = self.upsample(output)
+        upsample_outputs = upsample_outputs.permute(0, 2, 3, 1)  # BCHW -> BHWC
+
+        output_img_list = []
+        for upsample_output, meta in zip(upsample_outputs, meta_list):
+            c, s, w, h = meta['center'], meta['scale'], meta['width'], meta['height']
+            logits_result = transform_logits(upsample_output.data.cpu().numpy(), c, s, w, h, input_size=self.input_size)
+            parsing_result = np.argmax(logits_result, axis=2)
+            output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8))
+            output_img.putpalette(self.palette)
+            output_img_list.append(output_img)
+
+        return output_img_list[0] if len(output_img_list) == 1 else output_img_list

SCHP/networks/AugmentCE2P.py

@@ -0,0 +1,362 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+"""
+@Author  :   Peike Li
+@Contact :   peike.li@yahoo.com
+@File    :   AugmentCE2P.py
+@Time    :   8/4/19 3:35 PM
+@Desc    :
+@License :   This source code is licensed under the license found in the
+             LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from torch.nn import BatchNorm2d, LeakyReLU
+
+affine_par = True
+pretrained_settings = {
+    'resnet101': {
+        'imagenet': {
+            'input_space': 'BGR',
+            'input_size': [3, 224, 224],
+            'input_range': [0, 1],
+            'mean': [0.406, 0.456, 0.485],
+            'std': [0.225, 0.224, 0.229],
+            'num_classes': 1000
+        }
+    },
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, fist_dilation=1, multi_grid=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=dilation * multi_grid, dilation=dilation * multi_grid, bias=False)
+        self.bn2 = BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=False)
+        self.relu_inplace = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.dilation = dilation
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out = out + residual
+        out = self.relu_inplace(out)
+
+        return out
+
+
+class PSPModule(nn.Module):
+    """
+    Reference:
+        Zhao, Hengshuang, et al. *"Pyramid scene parsing network."*
+    """
+
+    def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
+        super(PSPModule, self).__init__()
+
+        self.stages = []
+        self.stages = nn.ModuleList([self._make_stage(features, out_features, size) for size in sizes])
+        self.bottleneck = nn.Sequential(
+            nn.Conv2d(features + len(sizes) * out_features, out_features, kernel_size=3, padding=1, dilation=1,
+                      bias=False),
+            BatchNorm2d(out_features),
+            LeakyReLU(),
+        )
+
+    def _make_stage(self, features, out_features, size):
+        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
+        conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
+        return nn.Sequential(
+            prior, 
+            conv, 
+            # bn
+            BatchNorm2d(out_features),
+            LeakyReLU(),
+        )
+
+    def forward(self, feats):
+        h, w = feats.size(2), feats.size(3)
+        priors = [F.interpolate(input=stage(feats), size=(h, w), mode='bilinear', align_corners=True) for stage in
+                  self.stages] + [feats]
+        bottle = self.bottleneck(torch.cat(priors, 1))
+        return bottle
+
+
+class ASPPModule(nn.Module):
+    """
+    Reference: 
+        Chen, Liang-Chieh, et al. *"Rethinking Atrous Convolution for Semantic Image Segmentation."*
+    """
+
+    def __init__(self, features, inner_features=256, out_features=512, dilations=(12, 24, 36)):
+        super(ASPPModule, self).__init__()
+
+        self.conv1 = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
+                                   nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1,
+                                             bias=False),
+                                #    InPlaceABNSync(inner_features)
+                                    BatchNorm2d(inner_features),
+                                    LeakyReLU(),
+                                )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=False),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            )
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[1], dilation=dilations[1], bias=False),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            )
+        self.conv5 = nn.Sequential(
+            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[2], dilation=dilations[2], bias=False),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            )
+
+        self.bottleneck = nn.Sequential(
+            nn.Conv2d(inner_features * 5, out_features, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            nn.Dropout2d(0.1)
+        )
+
+    def forward(self, x):
+        _, _, h, w = x.size()
+
+        feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True)
+
+        feat2 = self.conv2(x)
+        feat3 = self.conv3(x)
+        feat4 = self.conv4(x)
+        feat5 = self.conv5(x)
+        out = torch.cat((feat1, feat2, feat3, feat4, feat5), 1)
+
+        bottle = self.bottleneck(out)
+        return bottle
+
+
+class Edge_Module(nn.Module):
+    """
+    Edge Learning Branch
+    """
+
+    def __init__(self, in_fea=[256, 512, 1024], mid_fea=256, out_fea=2):
+        super(Edge_Module, self).__init__()
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_fea[0], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_fea[1], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(in_fea[2], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
+        )
+        self.conv4 = nn.Conv2d(mid_fea, out_fea, kernel_size=3, padding=1, dilation=1, bias=True)
+        # self.conv5 = nn.Conv2d(out_fea * 3, out_fea, kernel_size=1, padding=0, dilation=1, bias=True)
+
+    def forward(self, x1, x2, x3):
+        _, _, h, w = x1.size()
+
+        edge1_fea = self.conv1(x1)
+        # edge1 = self.conv4(edge1_fea)
+        edge2_fea = self.conv2(x2)
+        edge2 = self.conv4(edge2_fea)
+        edge3_fea = self.conv3(x3)
+        edge3 = self.conv4(edge3_fea)
+
+        edge2_fea = F.interpolate(edge2_fea, size=(h, w), mode='bilinear', align_corners=True)
+        edge3_fea = F.interpolate(edge3_fea, size=(h, w), mode='bilinear', align_corners=True)
+        edge2 = F.interpolate(edge2, size=(h, w), mode='bilinear', align_corners=True)
+        edge3 = F.interpolate(edge3, size=(h, w), mode='bilinear', align_corners=True)
+
+        # edge = torch.cat([edge1, edge2, edge3], dim=1)
+        edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
+        # edge = self.conv5(edge)
+
+        # return edge, edge_fea
+        return edge_fea
+
+
+class Decoder_Module(nn.Module):
+    """
+    Parsing Branch Decoder Module.
+    """
+
+    def __init__(self, num_classes):
+        super(Decoder_Module, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(512, 256, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(256),
+            LeakyReLU(),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(256, 48, kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(48),
+            LeakyReLU(),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(304, 256, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(256),
+            LeakyReLU(),
+            nn.Conv2d(256, 256, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(256),
+            LeakyReLU(),
+        )
+
+        # self.conv4 = nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
+
+    def forward(self, xt, xl):
+        _, _, h, w = xl.size()
+        xt = F.interpolate(self.conv1(xt), size=(h, w), mode='bilinear', align_corners=True)
+        xl = self.conv2(xl)
+        x = torch.cat([xt, xl], dim=1)
+        x = self.conv3(x)
+        # seg = self.conv4(x)
+        # return seg, x
+        return x
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, num_classes):
+        self.inplanes = 128
+        super(ResNet, self).__init__()
+        self.conv1 = conv3x3(3, 64, stride=2)
+        self.bn1 = BatchNorm2d(64)
+        self.relu1 = nn.ReLU(inplace=False)
+        self.conv2 = conv3x3(64, 64)
+        self.bn2 = BatchNorm2d(64)
+        self.relu2 = nn.ReLU(inplace=False)
+        self.conv3 = conv3x3(64, 128)
+        self.bn3 = BatchNorm2d(128)
+        self.relu3 = nn.ReLU(inplace=False)
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=2, multi_grid=(1, 1, 1))
+
+        self.context_encoding = PSPModule(2048, 512)
+
+        self.edge = Edge_Module()
+        self.decoder = Decoder_Module(num_classes)
+
+        self.fushion = nn.Sequential(
+            nn.Conv2d(1024, 256, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(256),
+            LeakyReLU(),
+            nn.Dropout2d(0.1),
+            nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
+        )
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, multi_grid=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm2d(planes * block.expansion, affine=affine_par))
+
+        layers = []
+        generate_multi_grid = lambda index, grids: grids[index % len(grids)] if isinstance(grids, tuple) else 1
+        layers.append(block(self.inplanes, planes, stride, dilation=dilation, downsample=downsample,
+                            multi_grid=generate_multi_grid(0, multi_grid)))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(
+                block(self.inplanes, planes, dilation=dilation, multi_grid=generate_multi_grid(i, multi_grid)))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.relu1(self.bn1(self.conv1(x)))
+        x = self.relu2(self.bn2(self.conv2(x)))
+        x = self.relu3(self.bn3(self.conv3(x)))
+        x = self.maxpool(x)
+        x2 = self.layer1(x)
+        x3 = self.layer2(x2)
+        x4 = self.layer3(x3)
+        x5 = self.layer4(x4)
+        x = self.context_encoding(x5)
+        # parsing_result, parsing_fea = self.decoder(x, x2)
+        parsing_fea = self.decoder(x, x2)
+        # Edge Branch
+        # edge_result, edge_fea = self.edge(x2, x3, x4)
+        edge_fea = self.edge(x2, x3, x4)
+        # Fusion Branch
+        x = torch.cat([parsing_fea, edge_fea], dim=1)
+        fusion_result = self.fushion(x)
+        # return [[parsing_result, fusion_result], [edge_result]]
+        return fusion_result
+
+
+def initialize_pretrained_model(model, settings, pretrained='./models/resnet101-imagenet.pth'):
+    model.input_space = settings['input_space']
+    model.input_size = settings['input_size']
+    model.input_range = settings['input_range']
+    model.mean = settings['mean']
+    model.std = settings['std']
+
+    if pretrained is not None:
+        saved_state_dict = torch.load(pretrained)
+        new_params = model.state_dict().copy()
+        for i in saved_state_dict:
+            i_parts = i.split('.')
+            if not i_parts[0] == 'fc':
+                new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
+        model.load_state_dict(new_params)
+
+
+def resnet101(num_classes=20, pretrained='./models/resnet101-imagenet.pth'):
+    model = ResNet(Bottleneck, [3, 4, 23, 3], num_classes)
+    settings = pretrained_settings['resnet101']['imagenet']
+    initialize_pretrained_model(model, settings, pretrained)
+    return model

SCHP/networks/__init__.py

@@ -0,0 +1,13 @@
+from __future__ import absolute_import
+
+from model.SCHP.networks.AugmentCE2P import resnet101
+
+__factory = {
+    'resnet101': resnet101,
+}
+
+
+def init_model(name, *args, **kwargs):
+    if name not in __factory.keys():
+        raise KeyError("Unknown model arch: {}".format(name))
+    return __factory[name](*args, **kwargs)

SCHP/utils/transforms.py

@@ -0,0 +1,167 @@
+# ------------------------------------------------------------------------------
+# Copyright (c) Microsoft
+# Licensed under the MIT License.
+# Written by Bin Xiao (Bin.Xiao@microsoft.com)
+# ------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import cv2
+import torch
+
+class BRG2Tensor_transform(object):
+    def __call__(self, pic):
+        img = torch.from_numpy(pic.transpose((2, 0, 1)))
+        if isinstance(img, torch.ByteTensor):
+            return img.float()
+        else:
+            return img
+
+class BGR2RGB_transform(object):
+    def __call__(self, tensor):
+        return tensor[[2,1,0],:,:]
+
+def flip_back(output_flipped, matched_parts):
+    '''
+    ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
+    '''
+    assert output_flipped.ndim == 4,\
+        'output_flipped should be [batch_size, num_joints, height, width]'
+
+    output_flipped = output_flipped[:, :, :, ::-1]
+
+    for pair in matched_parts:
+        tmp = output_flipped[:, pair[0], :, :].copy()
+        output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
+        output_flipped[:, pair[1], :, :] = tmp
+
+    return output_flipped
+
+
+def fliplr_joints(joints, joints_vis, width, matched_parts):
+    """
+    flip coords
+    """
+    # Flip horizontal
+    joints[:, 0] = width - joints[:, 0] - 1
+
+    # Change left-right parts
+    for pair in matched_parts:
+        joints[pair[0], :], joints[pair[1], :] = \
+            joints[pair[1], :], joints[pair[0], :].copy()
+        joints_vis[pair[0], :], joints_vis[pair[1], :] = \
+            joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
+
+    return joints*joints_vis, joints_vis
+
+
+def transform_preds(coords, center, scale, input_size):
+    target_coords = np.zeros(coords.shape)
+    trans = get_affine_transform(center, scale, 0, input_size, inv=1)
+    for p in range(coords.shape[0]):
+        target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
+    return target_coords
+
+def transform_parsing(pred, center, scale, width, height, input_size):
+
+    trans = get_affine_transform(center, scale, 0, input_size, inv=1)
+    target_pred = cv2.warpAffine(
+            pred,
+            trans,
+            (int(width), int(height)), #(int(width), int(height)),
+            flags=cv2.INTER_NEAREST,
+            borderMode=cv2.BORDER_CONSTANT,
+            borderValue=(0))
+
+    return target_pred
+
+def transform_logits(logits, center, scale, width, height, input_size):
+
+    trans = get_affine_transform(center, scale, 0, input_size, inv=1)
+    channel = logits.shape[2]
+    target_logits = []
+    for i in range(channel):
+        target_logit = cv2.warpAffine(
+            logits[:,:,i],
+            trans,
+            (int(width), int(height)), #(int(width), int(height)),
+            flags=cv2.INTER_LINEAR,
+            borderMode=cv2.BORDER_CONSTANT,
+            borderValue=(0))
+        target_logits.append(target_logit)
+    target_logits = np.stack(target_logits,axis=2)
+
+    return target_logits
+
+
+def get_affine_transform(center,
+                         scale,
+                         rot,
+                         output_size,
+                         shift=np.array([0, 0], dtype=np.float32),
+                         inv=0):
+    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
+        print(scale)
+        scale = np.array([scale, scale])
+
+    scale_tmp = scale
+
+    src_w = scale_tmp[0]
+    dst_w = output_size[1]
+    dst_h = output_size[0]
+
+    rot_rad = np.pi * rot / 180
+    src_dir = get_dir([0, src_w * -0.5], rot_rad)
+    dst_dir = np.array([0, (dst_w-1) * -0.5], np.float32)
+
+    src = np.zeros((3, 2), dtype=np.float32)
+    dst = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = center + scale_tmp * shift
+    src[1, :] = center + src_dir + scale_tmp * shift
+    dst[0, :] = [(dst_w-1) * 0.5, (dst_h-1) * 0.5]
+    dst[1, :] = np.array([(dst_w-1) * 0.5, (dst_h-1) * 0.5]) + dst_dir
+
+    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
+    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
+
+    if inv:
+        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    return trans
+
+
+def affine_transform(pt, t):
+    new_pt = np.array([pt[0], pt[1], 1.]).T
+    new_pt = np.dot(t, new_pt)
+    return new_pt[:2]
+
+
+def get_3rd_point(a, b):
+    direct = a - b
+    return b + np.array([-direct[1], direct[0]], dtype=np.float32)
+
+
+def get_dir(src_point, rot_rad):
+    sn, cs = np.sin(rot_rad), np.cos(rot_rad)
+
+    src_result = [0, 0]
+    src_result[0] = src_point[0] * cs - src_point[1] * sn
+    src_result[1] = src_point[0] * sn + src_point[1] * cs
+
+    return src_result
+
+
+def crop(img, center, scale, output_size, rot=0):
+    trans = get_affine_transform(center, scale, rot, output_size)
+
+    dst_img = cv2.warpAffine(img,
+                             trans,
+                             (int(output_size[1]), int(output_size[0])),
+                             flags=cv2.INTER_LINEAR)
+
+    return dst_img

attn_processor.py

@@ -0,0 +1,113 @@
+from torch.nn import functional as F
+import torch
+# from flash_attn import flash_attn_func
+
+class SkipAttnProcessor(torch.nn.Module):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        return hidden_states
+
+class AttnProcessor2_0(torch.nn.Module):
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(
+        self,
+        hidden_size=None,
+        cross_attention_dim=None,
+        **kwargs
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+        *args,
+        **kwargs,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+        # hidden_states = flash_attn_func(
+        #     query, key, value, dropout_p=0.0, causal=False
+        # )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+   

cloth_masker.py

@@ -0,0 +1,273 @@
+import os
+from PIL import Image
+from typing import Union
+import numpy as np
+import cv2
+from diffusers.image_processor import VaeImageProcessor
+import torch
+
+from model.SCHP import SCHP  # type: ignore
+from model.DensePose import DensePose  # type: ignore
+
+DENSE_INDEX_MAP = {
+    "background": [0],
+    "torso": [1, 2],
+    "right hand": [3],
+    "left hand": [4],
+    "right foot": [5],
+    "left foot": [6],
+    "right thigh": [7, 9],
+    "left thigh": [8, 10],
+    "right leg": [11, 13],
+    "left leg": [12, 14],
+    "left big arm": [15, 17],
+    "right big arm": [16, 18],
+    "left forearm": [19, 21],
+    "right forearm": [20, 22],
+    "face": [23, 24],
+    "thighs": [7, 8, 9, 10],
+    "legs": [11, 12, 13, 14],
+    "hands": [3, 4],
+    "feet": [5, 6],
+    "big arms": [15, 16, 17, 18],
+    "forearms": [19, 20, 21, 22],
+}
+
+ATR_MAPPING = {
+    'Background': 0, 'Hat': 1, 'Hair': 2, 'Sunglasses': 3, 
+    'Upper-clothes': 4, 'Skirt': 5, 'Pants': 6, 'Dress': 7,
+    'Belt': 8, 'Left-shoe': 9, 'Right-shoe': 10, 'Face': 11, 
+    'Left-leg': 12, 'Right-leg': 13, 'Left-arm': 14, 'Right-arm': 15,
+    'Bag': 16, 'Scarf': 17
+}
+
+LIP_MAPPING = {
+    'Background': 0, 'Hat': 1, 'Hair': 2, 'Glove': 3, 
+    'Sunglasses': 4, 'Upper-clothes': 5, 'Dress': 6, 'Coat': 7,
+    'Socks': 8, 'Pants': 9, 'Jumpsuits': 10, 'Scarf': 11, 
+    'Skirt': 12, 'Face': 13, 'Left-arm': 14, 'Right-arm': 15,
+    'Left-leg': 16, 'Right-leg': 17, 'Left-shoe': 18, 'Right-shoe': 19
+}
+
+PROTECT_BODY_PARTS = {
+    'upper': ['Left-leg', 'Right-leg'],
+    'lower': ['Right-arm', 'Left-arm', 'Face'],
+    'overall': [],
+    'inner': ['Left-leg', 'Right-leg'],
+    'outer': ['Left-leg', 'Right-leg'],
+}
+PROTECT_CLOTH_PARTS = {
+    'upper': {
+        'ATR': ['Skirt', 'Pants'],
+        'LIP': ['Skirt', 'Pants']
+    },
+    'lower': {
+        'ATR': ['Upper-clothes'],
+        'LIP': ['Upper-clothes', 'Coat']
+    },
+    'overall': {
+        'ATR': [],
+        'LIP': []
+    },
+    'inner': {
+        'ATR': ['Dress', 'Coat', 'Skirt', 'Pants'],
+        'LIP': ['Dress', 'Coat', 'Skirt', 'Pants', 'Jumpsuits']
+    },
+    'outer': {
+        'ATR': ['Dress', 'Pants', 'Skirt'],
+        'LIP': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Jumpsuits']
+    }
+}
+MASK_CLOTH_PARTS = {
+    'upper': ['Upper-clothes', 'Coat', 'Dress', 'Jumpsuits'],
+    'lower': ['Pants', 'Skirt', 'Dress', 'Jumpsuits'],
+    'overall': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Coat', 'Jumpsuits'],
+    'inner': ['Upper-clothes'],
+    'outer': ['Coat',]
+}
+MASK_DENSE_PARTS = {
+    'upper': ['torso', 'big arms', 'forearms'],
+    'lower': ['thighs', 'legs'],
+    'overall': ['torso', 'thighs', 'legs', 'big arms', 'forearms'],
+    'inner': ['torso'],
+    'outer': ['torso', 'big arms', 'forearms']
+}
+    
+schp_public_protect_parts = ['Hat', 'Hair', 'Sunglasses', 'Left-shoe', 'Right-shoe', 'Bag', 'Glove', 'Scarf']
+schp_protect_parts = {
+    'upper': ['Left-leg', 'Right-leg', 'Skirt', 'Pants', 'Jumpsuits'],  
+    'lower': ['Left-arm', 'Right-arm', 'Upper-clothes', 'Coat'],
+    'overall': [],
+    'inner': ['Left-leg', 'Right-leg', 'Skirt', 'Pants', 'Jumpsuits', 'Coat'],
+    'outer': ['Left-leg', 'Right-leg', 'Skirt', 'Pants', 'Jumpsuits', 'Upper-clothes']
+}
+schp_mask_parts = {
+    'upper': ['Upper-clothes', 'Dress', 'Coat', 'Jumpsuits'],
+    'lower': ['Pants', 'Skirt', 'Dress', 'Jumpsuits', 'socks'],
+    'overall': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Coat', 'Jumpsuits', 'socks'],
+    'inner': ['Upper-clothes'],
+    'outer': ['Coat',]
+}
+
+dense_mask_parts = {
+    'upper': ['torso', 'big arms', 'forearms'],
+    'lower': ['thighs', 'legs'],
+    'overall': ['torso', 'thighs', 'legs', 'big arms', 'forearms'],
+    'inner': ['torso'],
+    'outer': ['torso', 'big arms', 'forearms']
+}
+
+def vis_mask(image, mask):
+    image = np.array(image).astype(np.uint8)
+    mask = np.array(mask).astype(np.uint8)
+    mask[mask > 127] = 255
+    mask[mask <= 127] = 0
+    mask = np.expand_dims(mask, axis=-1)
+    mask = np.repeat(mask, 3, axis=-1)
+    mask = mask / 255
+    return Image.fromarray((image * (1 - mask)).astype(np.uint8))
+
+def part_mask_of(part: Union[str, list],
+                 parse: np.ndarray, mapping: dict):
+    if isinstance(part, str):
+        part = [part]
+    mask = np.zeros_like(parse)
+    for _ in part:
+        if _ not in mapping:
+            continue
+        if isinstance(mapping[_], list):
+            for i in mapping[_]:
+                mask += (parse == i)
+        else:
+            mask += (parse == mapping[_])
+    return mask
+
+def hull_mask(mask_area: np.ndarray):
+    ret, binary = cv2.threshold(mask_area, 127, 255, cv2.THRESH_BINARY)
+    contours, hierarchy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    hull_mask = np.zeros_like(mask_area)
+    for c in contours:
+        hull = cv2.convexHull(c)
+        hull_mask = cv2.fillPoly(np.zeros_like(mask_area), [hull], 255) | hull_mask
+    return hull_mask
+    
+
+class AutoMasker:
+    def __init__(
+        self, 
+        densepose_ckpt='./Models/DensePose', 
+        schp_ckpt='./Models/SCHP', 
+        device='cuda'):
+        np.random.seed(0)
+        torch.manual_seed(0)
+        torch.cuda.manual_seed(0)
+        
+        self.densepose_processor = DensePose(densepose_ckpt, device)
+        self.schp_processor_atr = SCHP(ckpt_path=os.path.join(schp_ckpt, 'exp-schp-201908301523-atr.pth'), device=device)
+        self.schp_processor_lip = SCHP(ckpt_path=os.path.join(schp_ckpt, 'exp-schp-201908261155-lip.pth'), device=device)
+        
+        self.mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True)
+
+    def process_densepose(self, image_or_path):
+        return self.densepose_processor(image_or_path, resize=1024)
+
+    def process_schp_lip(self, image_or_path):
+        return self.schp_processor_lip(image_or_path)
+
+    def process_schp_atr(self, image_or_path):
+        return self.schp_processor_atr(image_or_path)
+        
+    def preprocess_image(self, image_or_path):
+        return {
+            'densepose': self.densepose_processor(image_or_path, resize=1024),
+            'schp_atr': self.schp_processor_atr(image_or_path),
+            'schp_lip': self.schp_processor_lip(image_or_path)
+        }
+    
+    @staticmethod
+    def cloth_agnostic_mask(
+        densepose_mask: Image.Image,
+        schp_lip_mask: Image.Image,
+        schp_atr_mask: Image.Image,
+        part: str='overall',
+        **kwargs
+    ):
+        assert part in ['upper', 'lower', 'overall', 'inner', 'outer'], f"part should be one of ['upper', 'lower', 'overall', 'inner', 'outer'], but got {part}"
+        w, h = densepose_mask.size
+        
+        dilate_kernel = max(w, h) // 250
+        dilate_kernel = dilate_kernel if dilate_kernel % 2 == 1 else dilate_kernel + 1
+        dilate_kernel = np.ones((dilate_kernel, dilate_kernel), np.uint8)
+        
+        kernal_size = max(w, h) // 25
+        kernal_size = kernal_size if kernal_size % 2 == 1 else kernal_size + 1
+        
+        densepose_mask = np.array(densepose_mask)
+        schp_lip_mask = np.array(schp_lip_mask)
+        schp_atr_mask = np.array(schp_atr_mask)
+        
+        # Strong Protect Area (Hands, Face, Accessory, Feet)
+        hands_protect_area = part_mask_of(['hands', 'feet'], densepose_mask, DENSE_INDEX_MAP)
+        hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1)
+        hands_protect_area = hands_protect_area & \
+            (part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_atr_mask, ATR_MAPPING) | \
+             part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_lip_mask, LIP_MAPPING))
+        face_protect_area = part_mask_of('Face', schp_lip_mask, LIP_MAPPING)
+
+        strong_protect_area = hands_protect_area | face_protect_area 
+
+        # Weak Protect Area (Hair, Irrelevant Clothes, Body Parts)
+        body_protect_area = part_mask_of(PROTECT_BODY_PARTS[part], schp_lip_mask, LIP_MAPPING) | part_mask_of(PROTECT_BODY_PARTS[part], schp_atr_mask, ATR_MAPPING)
+        hair_protect_area = part_mask_of(['Hair'], schp_lip_mask, LIP_MAPPING) | \
+            part_mask_of(['Hair'], schp_atr_mask, ATR_MAPPING)
+        cloth_protect_area = part_mask_of(PROTECT_CLOTH_PARTS[part]['LIP'], schp_lip_mask, LIP_MAPPING) | \
+            part_mask_of(PROTECT_CLOTH_PARTS[part]['ATR'], schp_atr_mask, ATR_MAPPING)
+        accessory_protect_area = part_mask_of((accessory_parts := ['Hat', 'Glove', 'Sunglasses', 'Bag', 'Left-shoe', 'Right-shoe', 'Scarf', 'Socks']), schp_lip_mask, LIP_MAPPING) | \
+            part_mask_of(accessory_parts, schp_atr_mask, ATR_MAPPING) 
+        weak_protect_area = body_protect_area | cloth_protect_area | hair_protect_area | strong_protect_area | accessory_protect_area
+        
+        # Mask Area
+        strong_mask_area = part_mask_of(MASK_CLOTH_PARTS[part], schp_lip_mask, LIP_MAPPING) | \
+            part_mask_of(MASK_CLOTH_PARTS[part], schp_atr_mask, ATR_MAPPING)
+        background_area = part_mask_of(['Background'], schp_lip_mask, LIP_MAPPING) & part_mask_of(['Background'], schp_atr_mask, ATR_MAPPING)
+        mask_dense_area = part_mask_of(MASK_DENSE_PARTS[part], densepose_mask, DENSE_INDEX_MAP)
+        mask_dense_area = cv2.resize(mask_dense_area.astype(np.uint8), None, fx=0.25, fy=0.25, interpolation=cv2.INTER_NEAREST)
+        mask_dense_area = cv2.dilate(mask_dense_area, dilate_kernel, iterations=2)
+        mask_dense_area = cv2.resize(mask_dense_area.astype(np.uint8), None, fx=4, fy=4, interpolation=cv2.INTER_NEAREST)
+
+
+        mask_area = (np.ones_like(densepose_mask) & (~weak_protect_area) & (~background_area)) | mask_dense_area
+
+        mask_area = hull_mask(mask_area * 255) // 255  # Convex Hull to expand the mask area
+        mask_area = mask_area & (~weak_protect_area)
+        mask_area = cv2.GaussianBlur(mask_area * 255, (kernal_size, kernal_size), 0)
+        mask_area[mask_area < 25] = 0
+        mask_area[mask_area >= 25] = 1
+        mask_area = (mask_area | strong_mask_area) & (~strong_protect_area) 
+        mask_area = cv2.dilate(mask_area, dilate_kernel, iterations=1)
+
+        return Image.fromarray(mask_area * 255)
+        
+    def __call__(
+        self,
+        image: Union[str, Image.Image],
+        mask_type: str = "upper",
+    ):
+        assert mask_type in ['upper', 'lower', 'overall', 'inner', 'outer'], f"mask_type should be one of ['upper', 'lower', 'overall', 'inner', 'outer'], but got {mask_type}"
+        preprocess_results = self.preprocess_image(image)
+        mask = self.cloth_agnostic_mask(
+            preprocess_results['densepose'], 
+            preprocess_results['schp_lip'], 
+            preprocess_results['schp_atr'], 
+            part=mask_type,
+        )
+        return {
+            'mask': mask,
+            'densepose': preprocess_results['densepose'],
+            'schp_lip': preprocess_results['schp_lip'],
+            'schp_atr': preprocess_results['schp_atr']
+        }
+
+
+if __name__ == '__main__':
+    pass

flux/pipeline_flux_tryon.py

@@ -0,0 +1,499 @@
+
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import (
+    FluxLoraLoaderMixin,
+    FromSingleFileMixin,
+    TextualInversionLoaderMixin,
+)
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.pipelines.flux.pipeline_flux_fill import (
+    calculate_shift,
+    retrieve_latents,
+    retrieve_timesteps,
+)
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+
+from model.flux.transformer_flux import FluxTransformer2DModel
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+# Modified from `diffusers.pipelines.flux.pipeline_flux_fill.FluxFillPipeline`
+class FluxTryOnPipeline(
+    DiffusionPipeline,
+    FluxLoraLoaderMixin,
+    FromSingleFileMixin,
+    TextualInversionLoaderMixin,
+):
+    model_cpu_offload_seq = "transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents"]
+    
+    def __init__(
+        self, 
+        vae: AutoencoderKL, 
+        scheduler: FlowMatchEulerDiscreteScheduler, 
+        transformer: FluxTransformer2DModel,
+    ):
+        super().__init__()
+        self.register_modules(
+            vae=vae,
+            scheduler=scheduler,
+            transformer=transformer,
+        )
+        
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
+        )
+        
+        # Flux latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.mask_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor * 2,
+            vae_latent_channels=self.vae.config.latent_channels,
+            do_normalize=False,
+            do_binarize=True,
+            do_convert_grayscale=True,
+        )
+        self.default_sample_size = 128
+        
+        self.transformer.remove_text_layers() # TryOnEdit: remove text layers
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None, **kwargs):
+        transformer = FluxTransformer2DModel.from_pretrained(pretrained_model_name_or_path, subfolder="transformer")
+        transformer.remove_text_layers()
+        vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+        return FluxTryOnPipeline(vae, scheduler, transformer)
+    
+    def prepare_mask_latents(
+        self,
+        mask,
+        masked_image,
+        batch_size,
+        num_channels_latents,
+        num_images_per_prompt,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+    ):
+        # 1. calculate the height and width of the latents
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        # 2. encode the masked image
+        if masked_image.shape[1] == num_channels_latents:
+            masked_image_latents = masked_image
+        else:
+            masked_image_latents = retrieve_latents(self.vae.encode(masked_image), generator=generator)
+
+        masked_image_latents = (masked_image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
+
+        # 3. duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
+        batch_size = batch_size * num_images_per_prompt
+        if mask.shape[0] < batch_size:
+            if not batch_size % mask.shape[0] == 0:
+                raise ValueError(
+                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
+                    f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"
+                    " of masks that you pass is divisible by the total requested batch size."
+                )
+            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
+        if masked_image_latents.shape[0] < batch_size:
+            if not batch_size % masked_image_latents.shape[0] == 0:
+                raise ValueError(
+                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"
+                    f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
+                    " Make sure the number of images that you pass is divisible by the total requested batch size."
+                )
+            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)
+
+        # 4. pack the masked_image_latents
+        # batch_size, num_channels_latents, height, width -> batch_size, height//2 * width//2 , num_channels_latents*4
+        masked_image_latents = self._pack_latents(
+            masked_image_latents,
+            batch_size,
+            num_channels_latents,
+            height,
+            width,
+        )
+
+        # 5.resize mask to latents shape we we concatenate the mask to the latents
+        mask = mask[:, 0, :, :]  # batch_size, 8 * height, 8 * width (mask has not been 8x compressed)
+        mask = mask.view(
+            batch_size, height, self.vae_scale_factor, width, self.vae_scale_factor
+        )  # batch_size, height, 8, width, 8
+        mask = mask.permute(0, 2, 4, 1, 3)  # batch_size, 8, 8, height, width
+        mask = mask.reshape(
+            batch_size, self.vae_scale_factor * self.vae_scale_factor, height, width
+        )  # batch_size, 8*8, height, width
+
+        # 6. pack the mask:
+        # batch_size, 64, height, width -> batch_size, height//2 * width//2 , 64*2*2
+        mask = self._pack_latents(
+            mask,
+            batch_size,
+            self.vae_scale_factor * self.vae_scale_factor,
+            height,
+            width,
+        )
+        mask = mask.to(device=device, dtype=dtype)
+
+        return mask, masked_image_latents
+
+    def check_inputs(
+        self,
+        height,
+        width,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+        image=None,
+        mask_image=None,
+        condition_image=None,
+        masked_image_latents=None,
+    ):
+        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}")
+
+        if image is not None and masked_image_latents is not None:
+            raise ValueError(
+                "Please provide either  `image` or `masked_image_latents`, `masked_image_latents` should not be passed."
+            )
+
+        if image is not None and mask_image is None:
+            raise ValueError("Please provide `mask_image` when passing `image`.")
+        
+        if condition_image is None:
+            raise ValueError("Please provide `condition_image`.")
+
+    @staticmethod
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._prepare_latent_image_ids
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    @staticmethod
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._pack_latents
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._unpack_latents
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (vae_scale_factor * 2))
+        width = 2 * (int(width) // (vae_scale_factor * 2))
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+        return latents
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_latents
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is not None:
+            latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+            return latents.to(device=device, dtype=dtype), latent_image_ids
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+
+        latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+
+        return latents, latent_image_ids
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+    
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Optional[torch.FloatTensor] = None,
+        condition_image: Optional[torch.FloatTensor] = None,  # TryOnEdit: condition image (garment)
+        mask_image: Optional[torch.FloatTensor] = None,
+        masked_image_latents: Optional[torch.FloatTensor] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 30.0,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ):
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+        
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            height,
+            width,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+            image=image,
+            mask_image=mask_image,
+            condition_image=condition_image,
+            masked_image_latents=masked_image_latents,
+        )
+        
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+        
+        # 2. Define call parameters
+        batch_size = 1
+        device = self._execution_device
+        dtype = self.transformer.dtype
+        
+        # 3. Prepare prompt embeddings
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        
+        # 4. Prepare latent variables
+        num_channels_latents = self.vae.config.latent_channels
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width * 2, # TryOnEdit: width * 2
+            dtype,
+            device,
+            generator,
+            latents,
+        )
+        
+        # 5. Prepare mask and masked image latents
+        if masked_image_latents is not None:
+            masked_image_latents = masked_image_latents.to(latents.device)
+        else:
+            image = self.image_processor.preprocess(image, height=height, width=width)
+            condition_image = self.image_processor.preprocess(condition_image, height=height, width=width)
+            mask_image = self.mask_processor.preprocess(mask_image, height=height, width=width)
+
+            masked_image = image * (1 - mask_image)
+            masked_image = masked_image.to(device=device, dtype=dtype)
+            
+            # TryOnEdit: Concat condition image to masked image
+            condition_image = condition_image.to(device=device, dtype=dtype)
+            masked_image = torch.cat((masked_image, condition_image), dim=-1)
+            mask_image = torch.cat((mask_image, torch.zeros_like(mask_image)), dim=-1)
+
+            height, width = image.shape[-2:]
+            mask, masked_image_latents = self.prepare_mask_latents(
+                mask_image,
+                masked_image,
+                batch_size,
+                num_channels_latents,
+                num_images_per_prompt,
+                height,
+                width * 2, # TryOnEdit: width * 2
+                dtype,
+                device,
+                generator,
+            )
+            masked_image_latents = torch.cat((masked_image_latents, mask), dim=-1)
+        
+        # 6. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+        
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+        
+        # 7. Denoising loop
+        pooled_prompt_embeds = torch.zeros([latents.shape[0], 768], device=device, dtype=dtype) # TryOnEdit: for now, we don't use pooled prompt embeddings
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+                
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+                noise_pred = self.transformer(
+                    hidden_states=torch.cat((latents, masked_image_latents), dim=2),
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=None,
+                    txt_ids=None,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        # 8. Post-process the image
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width * 2, self.vae_scale_factor) # TryOnEdit: width * 2
+            latents = latents.split(latents.shape[-1] // 2, dim=-1)[0]  # TryOnEdit: split along the last dimension
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return FluxPipelineOutput(images=image)

flux/transformer_flux.py

@@ -0,0 +1,672 @@
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_processor import (
+    Attention,
+    AttentionProcessor,
+    FusedFluxAttnProcessor2_0,
+)   
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+# Modified from `diffusers.models.attention_processor.FluxAttnProcessor2_0`
+class FluxAttnProcessor2_0:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        if encoder_hidden_states is not None:
+            # `context` projections.
+            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+        
+        # edited for try-on
+        if not attn.pre_only:
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+
+        if encoder_hidden_states is not None:
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+
+@maybe_allow_in_graph
+class FluxSingleTransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
+        super().__init__()
+        self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+        self.norm = AdaLayerNormZeroSingle(dim)
+        self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+
+        processor = FluxAttnProcessor2_0()
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=processor,
+            qk_norm="rms_norm",
+            eps=1e-6,
+            pre_only=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        image_rotary_emb=None,
+        joint_attention_kwargs=None,
+    ):
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = residual + hidden_states
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return hidden_states
+
+
+@maybe_allow_in_graph
+class FluxTransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(self, dim, num_attention_heads, attention_head_dim, qk_norm="rms_norm", eps=1e-6):
+        super().__init__()
+
+        self.norm1 = AdaLayerNormZero(dim)
+
+        self.norm1_context = AdaLayerNormZero(dim)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            processor = FluxAttnProcessor2_0()
+        else:
+            raise ValueError(
+                "The current PyTorch version does not support the `scaled_dot_product_attention` function."
+            )
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=False,
+            bias=True,
+            processor=processor,
+            qk_norm=qk_norm,
+            eps=eps,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def remove_text_layers(self):
+        # for try-on, we don't need the text conditioning
+        self.norm1_context = None
+        self.ff_context = None
+        self.norm2_context = None
+        self.attn.to_added_qkv = None
+        self.attn.norm_added_q = None
+        self.attn.norm_added_k = None
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        image_rotary_emb=None,
+        joint_attention_kwargs=None,
+    ):
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+
+        if encoder_hidden_states is not None:
+            norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+                encoder_hidden_states, emb=temb
+            )
+        else:
+            norm_encoder_hidden_states = None
+
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        # Attention.
+        
+        outputs = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+        if isinstance(outputs, tuple):
+            attn_output, context_attn_output = outputs
+        else:
+            attn_output = outputs
+
+        # Process attention outputs for the `hidden_states`.
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = hidden_states + ff_output
+
+        # Process attention outputs for the `encoder_hidden_states`.
+        if encoder_hidden_states is not None:
+            context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+            encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+            norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+            norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+            context_ff_output = self.ff_context(norm_encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+            if encoder_hidden_states.dtype == torch.float16:
+                encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class FluxTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    """
+    The Transformer model introduced in Flux.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Parameters:
+        patch_size (`int`): Patch size to turn the input data into small patches.
+        in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
+        num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
+        num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
+        attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
+        num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
+        joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
+        guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["FluxTransformerBlock", "FluxSingleTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        out_channels: Optional[int] = None,
+        num_layers: int = 19,
+        num_single_layers: int = 38,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 4096,
+        pooled_projection_dim: int = 768,
+        guidance_embeds: bool = False,
+        axes_dims_rope: Tuple[int] = (16, 56, 56),
+    ):
+        super().__init__()
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
+        self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedFluxAttnProcessor2_0
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedFluxAttnProcessor2_0())
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def remove_text_layers(self):
+        self.context_embedder = None
+        for transformer_block in self.transformer_blocks:
+            transformer_block.remove_text_layers()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        condition_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
+    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+                
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        guidance = guidance.to(hidden_states.dtype) * 1000 if guidance is not None else None
+            
+        temb = self.time_text_embed(timestep, pooled_projections) if guidance is None else self.time_text_embed(timestep, guidance, pooled_projections)
+        
+        if encoder_hidden_states is not None:
+            encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        ids = torch.cat((txt_ids, img_ids), dim=0) if txt_ids is not None else img_ids  # for try-on, we don't need txt_ids
+        image_rotary_emb = self.pos_embed(ids)
+
+        # MMDiT Blocks
+        for index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                result = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+                if isinstance(result, tuple):
+                    encoder_hidden_states, hidden_states = result
+                else:
+                    hidden_states = result
+
+            else:
+                result = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+                if isinstance(result, tuple):
+                    encoder_hidden_states, hidden_states = result
+                else:
+                    hidden_states = result
+                    
+            # Condition residual (for try-on pose conditioning)
+            if condition_hidden_states is not None and index_block == 0:
+                hidden_states = hidden_states + condition_hidden_states
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                        hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        if encoder_hidden_states is not None:       
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        # Single DiT Blocks
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        if encoder_hidden_states is not None:
+            hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

pipeline.py

@@ -0,0 +1,332 @@
+import inspect
+import os
+from typing import Union
+
+import PIL
+import numpy as np
+import torch
+import tqdm
+from accelerate import load_checkpoint_in_model
+from diffusers import AutoencoderKL, DDIMScheduler, UNet2DConditionModel
+from diffusers.pipelines.stable_diffusion.safety_checker import \
+    StableDiffusionSafetyChecker
+from diffusers.utils.torch_utils import randn_tensor
+from huggingface_hub import snapshot_download
+from transformers import CLIPImageProcessor
+
+from model.attn_processor import SkipAttnProcessor
+from model.utils import get_trainable_module, init_adapter
+from utils import (compute_vae_encodings, numpy_to_pil, prepare_image,
+                   prepare_mask_image, resize_and_crop, resize_and_padding)
+
+
+class CatVTONPipeline:
+    def __init__(
+        self, 
+        base_ckpt, 
+        attn_ckpt, 
+        attn_ckpt_version="mix",
+        weight_dtype=torch.float32,
+        device='cuda',
+        compile=False,
+        skip_safety_check=False,
+        use_tf32=True,
+    ):
+        self.device = device
+        self.weight_dtype = weight_dtype
+        self.skip_safety_check = skip_safety_check
+
+        self.noise_scheduler = DDIMScheduler.from_pretrained(base_ckpt, subfolder="scheduler")
+        self.vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device, dtype=weight_dtype)
+        if not skip_safety_check:
+            self.feature_extractor = CLIPImageProcessor.from_pretrained(base_ckpt, subfolder="feature_extractor")
+            self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(base_ckpt, subfolder="safety_checker").to(device, dtype=weight_dtype)
+        self.unet = UNet2DConditionModel.from_pretrained(base_ckpt, subfolder="unet").to(device, dtype=weight_dtype)
+        init_adapter(self.unet, cross_attn_cls=SkipAttnProcessor)  # Skip Cross-Attention
+        self.attn_modules = get_trainable_module(self.unet, "attention")
+        self.auto_attn_ckpt_load(attn_ckpt, attn_ckpt_version)
+        # Pytorch 2.0 Compile
+        if compile:
+            self.unet = torch.compile(self.unet)
+            self.vae = torch.compile(self.vae, mode="reduce-overhead")
+            
+        # Enable TF32 for faster training on Ampere GPUs (A100 and RTX 30 series).
+        if use_tf32:
+            torch.set_float32_matmul_precision("high")
+            torch.backends.cuda.matmul.allow_tf32 = True
+
+    def auto_attn_ckpt_load(self, attn_ckpt, version):
+        sub_folder = {
+            "mix": "mix-48k-1024",
+            "vitonhd": "vitonhd-16k-512",
+            "dresscode": "dresscode-16k-512",
+        }[version]
+        if os.path.exists(attn_ckpt):
+            load_checkpoint_in_model(self.attn_modules, os.path.join(attn_ckpt, sub_folder, 'attention'))
+        else:
+            repo_path = snapshot_download(repo_id=attn_ckpt)
+            print(f"Downloaded {attn_ckpt} to {repo_path}")
+            load_checkpoint_in_model(self.attn_modules, os.path.join(repo_path, sub_folder, 'attention'))
+            
+    def run_safety_checker(self, image):
+        if self.safety_checker is None:
+            has_nsfw_concept = None
+        else:
+            safety_checker_input = self.feature_extractor(image, return_tensors="pt").to(self.device)
+            image, has_nsfw_concept = self.safety_checker(
+                images=image, clip_input=safety_checker_input.pixel_values.to(self.weight_dtype)
+            )
+        return image, has_nsfw_concept
+    
+    def check_inputs(self, image, condition_image, mask, width, height):
+        if isinstance(image, torch.Tensor) and isinstance(condition_image, torch.Tensor) and isinstance(mask, torch.Tensor):
+            return image, condition_image, mask
+        assert image.size == mask.size, "Image and mask must have the same size"
+        image = resize_and_crop(image, (width, height))
+        mask = resize_and_crop(mask, (width, height))
+        condition_image = resize_and_padding(condition_image, (width, height))
+        return image, condition_image, mask
+    
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(
+            inspect.signature(self.noise_scheduler.step).parameters.keys()
+        )
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(
+            inspect.signature(self.noise_scheduler.step).parameters.keys()
+        )
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    @torch.no_grad()
+    def __call__(
+        self, 
+        image: Union[PIL.Image.Image, torch.Tensor],
+        condition_image: Union[PIL.Image.Image, torch.Tensor],
+        mask: Union[PIL.Image.Image, torch.Tensor],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 2.5,
+        height: int = 1024,
+        width: int = 768,
+        generator=None,
+        eta=1.0,
+        **kwargs
+    ):
+        concat_dim = -2  # FIXME: y axis concat
+        # Prepare inputs to Tensor
+        image, condition_image, mask = self.check_inputs(image, condition_image, mask, width, height)
+        image = prepare_image(image).to(self.device, dtype=self.weight_dtype)
+        condition_image = prepare_image(condition_image).to(self.device, dtype=self.weight_dtype)
+        mask = prepare_mask_image(mask).to(self.device, dtype=self.weight_dtype)
+        # Mask image
+        masked_image = image * (mask < 0.5)
+        # VAE encoding
+        masked_latent = compute_vae_encodings(masked_image, self.vae)
+        condition_latent = compute_vae_encodings(condition_image, self.vae)
+        mask_latent = torch.nn.functional.interpolate(mask, size=masked_latent.shape[-2:], mode="nearest")
+        del image, mask, condition_image
+        # Concatenate latents
+        masked_latent_concat = torch.cat([masked_latent, condition_latent], dim=concat_dim)
+        mask_latent_concat = torch.cat([mask_latent, torch.zeros_like(mask_latent)], dim=concat_dim)
+        # Prepare noise
+        latents = randn_tensor(
+            masked_latent_concat.shape,
+            generator=generator,
+            device=masked_latent_concat.device,
+            dtype=self.weight_dtype,
+        )
+        # Prepare timesteps
+        self.noise_scheduler.set_timesteps(num_inference_steps, device=self.device)
+        timesteps = self.noise_scheduler.timesteps
+        latents = latents * self.noise_scheduler.init_noise_sigma
+        # Classifier-Free Guidance
+        if do_classifier_free_guidance := (guidance_scale > 1.0):
+            masked_latent_concat = torch.cat(
+                [
+                    torch.cat([masked_latent, torch.zeros_like(condition_latent)], dim=concat_dim),
+                    masked_latent_concat,
+                ]
+            )
+            mask_latent_concat = torch.cat([mask_latent_concat] * 2)
+
+        # Denoising loop
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_warmup_steps = (len(timesteps) - num_inference_steps * self.noise_scheduler.order)
+        with tqdm.tqdm(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                non_inpainting_latent_model_input = (torch.cat([latents] * 2) if do_classifier_free_guidance else latents)
+                non_inpainting_latent_model_input = self.noise_scheduler.scale_model_input(non_inpainting_latent_model_input, t)
+                # prepare the input for the inpainting model
+                inpainting_latent_model_input = torch.cat([non_inpainting_latent_model_input, mask_latent_concat, masked_latent_concat], dim=1)
+                # predict the noise residual
+                noise_pred= self.unet(
+                    inpainting_latent_model_input,
+                    t.to(self.device),
+                    encoder_hidden_states=None, # FIXME
+                    return_dict=False,
+                )[0]
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.noise_scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs
+                ).prev_sample
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps
+                    and (i + 1) % self.noise_scheduler.order == 0
+                ):
+                    progress_bar.update()
+
+        # Decode the final latents
+        latents = latents.split(latents.shape[concat_dim] // 2, dim=concat_dim)[0]
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents.to(self.device, dtype=self.weight_dtype)).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        image = numpy_to_pil(image)
+        
+        # Safety Check
+        if not self.skip_safety_check:
+            current_script_directory = os.path.dirname(os.path.realpath(__file__))
+            nsfw_image = os.path.join(os.path.dirname(current_script_directory), 'resource', 'img', 'NSFW.jpg')
+            nsfw_image = PIL.Image.open(nsfw_image).resize(image[0].size)
+            image_np = np.array(image)
+            _, has_nsfw_concept = self.run_safety_checker(image=image_np)
+            for i, not_safe in enumerate(has_nsfw_concept):
+                if not_safe:
+                    image[i] = nsfw_image
+        return image
+
+
+class CatVTONPix2PixPipeline(CatVTONPipeline):
+    def auto_attn_ckpt_load(self, attn_ckpt, version):
+        # TODO: Temperal fix for the model version
+        if os.path.exists(attn_ckpt):
+            load_checkpoint_in_model(self.attn_modules, os.path.join(attn_ckpt, version, 'attention'))
+        else:
+            repo_path = snapshot_download(repo_id=attn_ckpt)
+            print(f"Downloaded {attn_ckpt} to {repo_path}")
+            load_checkpoint_in_model(self.attn_modules, os.path.join(repo_path, version, 'attention'))
+    
+    def check_inputs(self, image, condition_image, width, height):
+        if isinstance(image, torch.Tensor) and isinstance(condition_image, torch.Tensor) and isinstance(torch.Tensor):
+            return image, condition_image
+        image = resize_and_crop(image, (width, height))
+        condition_image = resize_and_padding(condition_image, (width, height))
+        return image, condition_image
+
+    @torch.no_grad()
+    def __call__(
+        self, 
+        image: Union[PIL.Image.Image, torch.Tensor],
+        condition_image: Union[PIL.Image.Image, torch.Tensor],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 2.5,
+        height: int = 1024,
+        width: int = 768,
+        generator=None,
+        eta=1.0,
+        **kwargs
+    ):
+        concat_dim = -1
+        # Prepare inputs to Tensor
+        image, condition_image = self.check_inputs(image, condition_image, width, height)
+        image = prepare_image(image).to(self.device, dtype=self.weight_dtype)
+        condition_image = prepare_image(condition_image).to(self.device, dtype=self.weight_dtype)
+        # VAE encoding
+        image_latent = compute_vae_encodings(image, self.vae)
+        condition_latent = compute_vae_encodings(condition_image, self.vae)
+        del image, condition_image
+        # Concatenate latents
+        condition_latent_concat = torch.cat([image_latent, condition_latent], dim=concat_dim)
+        # Prepare noise
+        latents = randn_tensor(
+            condition_latent_concat.shape,
+            generator=generator,
+            device=condition_latent_concat.device,
+            dtype=self.weight_dtype,
+        )
+        # Prepare timesteps
+        self.noise_scheduler.set_timesteps(num_inference_steps, device=self.device)
+        timesteps = self.noise_scheduler.timesteps
+        latents = latents * self.noise_scheduler.init_noise_sigma
+        # Classifier-Free Guidance
+        if do_classifier_free_guidance := (guidance_scale > 1.0):
+            condition_latent_concat = torch.cat(
+                [
+                    torch.cat([image_latent, torch.zeros_like(condition_latent)], dim=concat_dim),
+                    condition_latent_concat,
+                ]
+            )
+
+        # Denoising loop
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_warmup_steps = (len(timesteps) - num_inference_steps * self.noise_scheduler.order)
+        with tqdm.tqdm(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = (torch.cat([latents] * 2) if do_classifier_free_guidance else latents)
+                latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, t)
+                # prepare the input for the inpainting model
+                p2p_latent_model_input = torch.cat([latent_model_input, condition_latent_concat], dim=1)
+                # predict the noise residual
+                noise_pred= self.unet(
+                    p2p_latent_model_input,
+                    t.to(self.device),
+                    encoder_hidden_states=None, 
+                    return_dict=False,
+                )[0]
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.noise_scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs
+                ).prev_sample
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps
+                    and (i + 1) % self.noise_scheduler.order == 0
+                ):
+                    progress_bar.update()
+
+        # Decode the final latents
+        latents = latents.split(latents.shape[concat_dim] // 2, dim=concat_dim)[0]
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents.to(self.device, dtype=self.weight_dtype)).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        image = numpy_to_pil(image)
+        
+        # Safety Check
+        if not self.skip_safety_check:
+            current_script_directory = os.path.dirname(os.path.realpath(__file__))
+            nsfw_image = os.path.join(os.path.dirname(current_script_directory), 'resource', 'img', 'NSFW.jpg')
+            nsfw_image = PIL.Image.open(nsfw_image).resize(image[0].size)
+            image_np = np.array(image)
+            _, has_nsfw_concept = self.run_safety_checker(image=image_np)
+            for i, not_safe in enumerate(has_nsfw_concept):
+                if not_safe:
+                    image[i] = nsfw_image
+        return image

utils.py

@@ -0,0 +1,84 @@
+import os
+import json
+import torch
+from model.attn_processor import AttnProcessor2_0, SkipAttnProcessor 
+
+
+def init_adapter(unet, 
+                 cross_attn_cls=SkipAttnProcessor,
+                 self_attn_cls=None,
+                 cross_attn_dim=None, 
+                 **kwargs):
+    if cross_attn_dim is None:
+        cross_attn_dim = unet.config.cross_attention_dim
+    attn_procs = {}
+    for name in unet.attn_processors.keys():
+        cross_attention_dim = None if name.endswith("attn1.processor") else cross_attn_dim
+        if name.startswith("mid_block"):
+            hidden_size = unet.config.block_out_channels[-1]
+        elif name.startswith("up_blocks"):
+            block_id = int(name[len("up_blocks.")])
+            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
+        elif name.startswith("down_blocks"):
+            block_id = int(name[len("down_blocks.")])
+            hidden_size = unet.config.block_out_channels[block_id]
+        if cross_attention_dim is None:
+            if self_attn_cls is not None:
+                attn_procs[name] = self_attn_cls(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, **kwargs)
+            else:
+                # retain the original attn processor
+                attn_procs[name] = AttnProcessor2_0(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, **kwargs)
+        else:
+            attn_procs[name] = cross_attn_cls(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, **kwargs)
+                                                    
+    unet.set_attn_processor(attn_procs)
+    adapter_modules = torch.nn.ModuleList(unet.attn_processors.values())
+    return adapter_modules
+
+def init_diffusion_model(diffusion_model_name_or_path, unet_class=None):
+    from diffusers import AutoencoderKL
+    from transformers import CLIPTextModel, CLIPTokenizer
+
+    text_encoder = CLIPTextModel.from_pretrained(diffusion_model_name_or_path, subfolder="text_encoder")
+    vae = AutoencoderKL.from_pretrained(diffusion_model_name_or_path, subfolder="vae")
+    tokenizer = CLIPTokenizer.from_pretrained(diffusion_model_name_or_path, subfolder="tokenizer")
+    try:
+        unet_folder = os.path.join(diffusion_model_name_or_path, "unet")
+        unet_configs = json.load(open(os.path.join(unet_folder, "config.json"), "r"))
+        unet = unet_class(**unet_configs)
+        unet.load_state_dict(torch.load(os.path.join(unet_folder, "diffusion_pytorch_model.bin"), map_location="cpu"), strict=True)
+    except:
+        unet = None
+    return text_encoder, vae, tokenizer, unet
+
+def attn_of_unet(unet):
+    attn_blocks = torch.nn.ModuleList()
+    for name, param in unet.named_modules():
+        if "attn1" in name:
+            attn_blocks.append(param)
+    return attn_blocks
+
+def get_trainable_module(unet, trainable_module_name):
+    if trainable_module_name == "unet":
+        return unet
+    elif trainable_module_name == "transformer":
+        trainable_modules = torch.nn.ModuleList()
+        for blocks in [unet.down_blocks, unet.mid_block, unet.up_blocks]:
+            if hasattr(blocks, "attentions"):
+                trainable_modules.append(blocks.attentions)
+            else:
+                for block in blocks:
+                    if hasattr(block, "attentions"):
+                        trainable_modules.append(block.attentions)
+        return trainable_modules
+    elif trainable_module_name == "attention":
+        attn_blocks = torch.nn.ModuleList()
+        for name, param in unet.named_modules():
+            if "attn1" in name:
+                attn_blocks.append(param)
+        return attn_blocks
+    else:
+        raise ValueError(f"Unknown trainable_module_name: {trainable_module_name}")
+
+                
+