libcom 0.1.8__py3-none-any.whl

libcom/shadow_generation/source/attention_processor.py

@@ -0,0 +1,234 @@
+# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import inspect
+import warnings
+from packaging import version
+
+class AttnProcessor(nn.Module):
+    r"""
+    Default processor for performing attention-related computations.
+    """
+
+    def __init__(
+        self,
+        hidden_size=None,
+        cross_attention_dim=None,
+    ):
+        super().__init__()
+
+    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
+        )
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+        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)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # 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
+
+class IPAttnProcessor(nn.Module):
+    r"""
+    Attention processor for IP-Adapater.
+    Args:
+        hidden_size (`int`):
+            The hidden size of the attention layer.
+        cross_attention_dim (`int`):
+            The number of channels in the `encoder_hidden_states`.
+        scale (`float`, defaults to 1.0):
+            the weight scale of image prompt.
+        num_tokens (`int`, defaults to 4 when do ip_adapter_plus it should be 16):
+            The context length of the image features.
+    """
+
+    def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=64):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.scale = scale
+        self.num_tokens = num_tokens
+
+        self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+        self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+
+    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
+        )
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+        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
+        else:
+            # get encoder_hidden_states, ip_hidden_states
+            end_pos = encoder_hidden_states.shape[1] - self.num_tokens
+            encoder_hidden_states, ip_hidden_states = (
+                encoder_hidden_states[:, :end_pos, :],
+                encoder_hidden_states[:, end_pos:, :],
+            )
+            if 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)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # for ip-adapter
+        ip_key = self.to_k_ip(ip_hidden_states)
+        ip_value = self.to_v_ip(ip_hidden_states)
+
+        ip_key = attn.head_to_batch_dim(ip_key)
+        ip_value = attn.head_to_batch_dim(ip_value)
+
+        ip_attention_probs = attn.get_attention_scores(query, ip_key, None)
+        self.attn_map = ip_attention_probs
+        ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
+        ip_hidden_states = attn.batch_to_head_dim(ip_hidden_states)
+
+        hidden_states = hidden_states + self.scale * ip_hidden_states
+
+        # 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
+    
+def deprecate(*args, take_from: Optional[Union[Dict, Any]] = None, standard_warn=True, stacklevel=2):
+    from .. import __version__
+
+    deprecated_kwargs = take_from
+    values = ()
+    if not isinstance(args[0], tuple):
+        args = (args,)
+
+    for attribute, version_name, message in args:
+        if version.parse(version.parse(__version__).base_version) >= version.parse(version_name):
+            raise ValueError(
+                f"The deprecation tuple {(attribute, version_name, message)} should be removed since diffusers'"
+                f" version {__version__} is >= {version_name}"
+            )
+
+        warning = None
+        if isinstance(deprecated_kwargs, dict) and attribute in deprecated_kwargs:
+            values += (deprecated_kwargs.pop(attribute),)
+            warning = f"The `{attribute}` argument is deprecated and will be removed in version {version_name}."
+        elif hasattr(deprecated_kwargs, attribute):
+            values += (getattr(deprecated_kwargs, attribute),)
+            warning = f"The `{attribute}` attribute is deprecated and will be removed in version {version_name}."
+        elif deprecated_kwargs is None:
+            warning = f"`{attribute}` is deprecated and will be removed in version {version_name}."
+
+        if warning is not None:
+            warning = warning + " " if standard_warn else ""
+            warnings.warn(warning + message, FutureWarning, stacklevel=stacklevel)
+
+    if isinstance(deprecated_kwargs, dict) and len(deprecated_kwargs) > 0:
+        call_frame = inspect.getouterframes(inspect.currentframe())[1]
+        filename = call_frame.filename
+        line_number = call_frame.lineno
+        function = call_frame.function
+        key, value = next(iter(deprecated_kwargs.items()))
+        raise TypeError(f"{function} in {filename} line {line_number - 1} got an unexpected keyword argument `{key}`")
+
+    if len(values) == 0:
+        return
+    elif len(values) == 1:
+        return values[0]
+    return values

libcom/shadow_generation/source/cldm/base_network.py

@@ -0,0 +1,166 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+class BaseNetwork(nn.Module):
+    def __init__(self):
+        super(BaseNetwork, self).__init__()
+
+    def init_weights(self, init_type='normal', gain=0.02):
+        def init_func(m):
+            classname = m.__class__.__name__
+            if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+                if init_type == 'normal':
+                    nn.init.normal_(m.weight.data, 0.0, gain)
+                elif init_type == 'xavier':
+                    nn.init.xavier_normal_(m.weight.data, gain=gain)
+                elif init_type == 'kaiming':
+                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+                elif init_type == 'orthogonal':
+                    nn.init.orthogonal_(m.weight.data, gain=gain)
+
+                if hasattr(m, 'bias') and m.bias is not None:
+                    nn.init.constant_(m.bias.data, 0.0)
+
+            elif classname.find('BatchNorm2d') != -1:
+                nn.init.normal_(m.weight.data, 1.0, gain)
+                nn.init.constant_(m.bias.data, 0.0)
+
+        self.apply(init_func)
+
+class DoubleConv(BaseNetwork):
+    """(convolution => [BN] => ReLU) * 2"""
+    def __init__(self, in_channels, out_channels, mid_channels=None, norm_layer = nn.InstanceNorm2d):
+        super().__init__()
+        self.norm_layer = norm_layer
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
+            self.norm_layer(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
+            self.norm_layer(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+class Bottleneck(BaseNetwork):
+    def __init__(self,in_places,out_places,stride=1,downsampling=False,expansion=4, norm_layer = nn.InstanceNorm2d):
+        super(Bottleneck,self).__init__()
+        self.expansion = expansion
+        self.downsampling = downsampling
+        self.norm_layer = norm_layer
+        self.bottleneck = nn.Sequential(
+            nn.Conv2d(in_channels=in_places,out_channels=out_places,kernel_size=1,stride=1),
+            self.norm_layer(out_places),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels=out_places,out_channels=out_places,kernel_size=3,stride=stride,padding=1),
+            self.norm_layer(out_places),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels=out_places,out_channels=out_places*self.expansion,kernel_size=1,stride=1),
+            self.norm_layer(out_places*self.expansion)
+        )
+
+        if self.downsampling :
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_channels=in_places,out_channels=out_places*self.expansion,kernel_size=1,stride=stride),
+                self.norm_layer(out_places*self.expansion)
+            )
+        self.relu = nn.ReLU(inplace=True)
+        
+    def forward(self,x):
+        residual = x
+        out = self.bottleneck(x)
+
+        if self.downsampling:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+        return out
+    
+class ResNet(BaseNetwork):
+    def __init__(self, blocks, out_places=[64,16,32,64,64], input_nc = 4, expansion=4, norm_layer=nn.InstanceNorm2d):
+        super(ResNet, self).__init__()
+        self.expansion = expansion
+        self.norm_layer = norm_layer
+        self.bilinear = True 
+        self.conv = DoubleConv(input_nc, 64, norm_layer=nn.BatchNorm2d)  #TODO:
+        self.conv.eval()
+
+        self.layer1 = self.make_layer(in_places=out_places[0], out_places=out_places[1], block=blocks[0], stride=2,
+                                      norm_layer=self.norm_layer)
+        self.layer2 = self.make_layer(in_places=4*out_places[1], out_places=out_places[2], block=blocks[1], stride=2,
+                                      norm_layer=self.norm_layer)
+        self.layer3 = self.make_layer(in_places=4*out_places[2], out_places=out_places[3], block=blocks[2], stride=2,
+                                      norm_layer=self.norm_layer)
+        self.layer4 = self.make_layer(in_places=4*out_places[3], out_places=out_places[4], block=blocks[3], stride=2,
+                                      norm_layer=self.norm_layer)
+
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+
+    def make_layer(self, in_places, out_places, block, stride, norm_layer=nn.InstanceNorm2d):
+        layers = []
+        layers.append(
+            Bottleneck(in_places, out_places, stride, downsampling=True, norm_layer=norm_layer))
+        for i in range(1, block):
+            layers.append(
+                Bottleneck(out_places * self.expansion, out_places, norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, input):
+        f1 = self.conv(input)
+        f2 = self.layer1(f1)
+        f3 = self.layer2(f2)
+        f4 = self.layer3(f3)
+        f5 = self.layer4(f4)
+        return f5
+
+class RegNetwork(BaseNetwork):
+
+    def __init__(self,input_dim=256, norm_layer=nn.InstanceNorm2d):
+        super(RegNetwork,self).__init__()
+        self.norm_layer = norm_layer
+        self.FeatureEncoder = ResNet([3,4,6,3],input_nc=4,norm_layer=self.norm_layer)
+        self.net = nn.Sequential(
+            nn.Conv2d(input_dim,input_dim*2,kernel_size=3,padding=1,stride=2),
+            self.norm_layer(input_dim*2),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(input_dim*2,input_dim*2,kernel_size=3,padding=1,stride=1),
+            self.norm_layer(input_dim*2),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(input_dim*2,input_dim*2,kernel_size=3,padding=1,stride=1),
+            self.norm_layer(input_dim*2),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+        )
+        self.fc_bbx = nn.Linear(input_dim*2, 5)
+
+    def forward(self, x):
+        x = self.FeatureEncoder(x)
+        x = self.net(x)
+        pred_t = self.fc_bbx(x)
+        pred_t[:, 4] = torch.tanh(pred_t[:, 4]) * (torch.pi / 2)
+        return pred_t
+
+class MaskCls(BaseNetwork):
+
+    def __init__(self, num_classes=256):
+        super(MaskCls, self).__init__()
+        # 加载预训练的ResNet模型，这里选择ResNet50作为基础模型
+        self.resnet = models.resnet50(pretrained=True)
+        self.resnet.conv1 = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        # 替换最后一层全连接层
+        num_features = self.resnet.fc.in_features
+        self.resnet.fc = nn.Linear(num_features, num_classes)
+    
+    def forward(self, x):
+        x = self.resnet(x)
+        return x