1230-image-colorization

submit new algorithm for image colorization and corresponding pipeline.
        Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11179627

data/test/images/audrey_hepburn.jpg

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da5cf2f3318e61cd38193af374b21a2dec0e90f2aa0e25b3b1825488eadbdc9d
+size 97191

modelscope/metainfo.py

@@ -67,6 +67,7 @@ class Models(object):
     real_basicvsr = 'real-basicvsr'
     rcp_sceneflow_estimation = 'rcp-sceneflow-estimation'
     image_casmvs_depth_estimation = 'image-casmvs-depth-estimation'
+    ddcolor = 'ddcolor'
 
     # EasyCV models
     yolox = 'YOLOX'
@@ -267,6 +268,7 @@ class Pipelines(object):
     video_super_resolution = 'realbasicvsr-video-super-resolution'
     pointcloud_sceneflow_estimation = 'pointcloud-sceneflow-estimation'
     image_multi_view_depth_estimation = 'image-multi-view-depth-estimation'
+    ddcolor_image_colorization = 'ddcolor-image-colorization'
 
     # nlp tasks
     automatic_post_editing = 'automatic-post-editing'

modelscope/models/cv/image_colorization/__init__.py

@@ -4,13 +4,13 @@ from typing import TYPE_CHECKING
 from modelscope.utils.import_utils import LazyImportModule
 
 if TYPE_CHECKING:
-    from .unet import DynamicUnetWide, DynamicUnetDeep
-    from .utils import NormType
+    from .unet import DynamicUnetWide, DynamicUnetDeep, NormType
+    from .ddcolor import DDColorForImageColorization
 
 else:
     _import_structure = {
-        'unet': ['DynamicUnetWide', 'DynamicUnetDeep'],
-        'utils': ['NormType']
+        'unet': ['DynamicUnetWide', 'DynamicUnetDeep', 'NormType'],
+        'ddcolor': ['DDColorForImageColorization'],
     }
 
     import sys

modelscope/models/cv/image_colorization/ddcolor/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .ddcolor_for_image_colorization import DDColorForImageColorization

modelscope/models/cv/image_colorization/ddcolor/ddcolor.py

@@ -0,0 +1,283 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import torch
+import torch.nn as nn
+
+from .utils.convnext import ConvNeXt
+from .utils.position_encoding import PositionEmbeddingSine
+from .utils.transformer_utils import (MLP, CrossAttentionLayer, FFNLayer,
+                                      SelfAttentionLayer)
+from .utils.unet import (CustomPixelShuffle_ICNR, Hook, NormType,
+                         UnetBlockWide, custom_conv_layer)
+
+
+class DDColor(nn.Module):
+
+    def __init__(self,
+                 encoder_name='convnext-l',
+                 input_size=(256, 256),
+                 num_queries=100):
+
+        super().__init__()
+
+        self.encoder = Encoder(encoder_name,
+                               ['norm0', 'norm1', 'norm2', 'norm3'])
+        self.encoder.eval()
+        test_input = torch.randn(1, 3, *input_size)
+        self.encoder(test_input)
+
+        self.decoder = Decoder(
+            self.encoder.hooks,
+            nf=512,
+            last_norm='Spectral',
+            num_queries=num_queries,
+            num_scales=3,
+            dec_layers=9,
+        )
+        self.refine_net = nn.Sequential(
+            custom_conv_layer(
+                num_queries + 3,
+                2,
+                ks=1,
+                use_activ=False,
+                norm_type=NormType.Spectral))
+
+        self.register_buffer(
+            'mean',
+            torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer(
+            'std',
+            torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+    def normalize(self, img):
+        return (img - self.mean) / self.std
+
+    def forward(self, img):
+        if img.shape[1] == 3:
+            img = self.normalize(img)
+
+        self.encoder(img)
+        out_feat = self.decoder()
+        coarse_input = torch.cat([out_feat, img], dim=1)
+        out = self.refine_net(coarse_input)
+
+        return out
+
+
+class Decoder(nn.Module):
+
+    def __init__(self,
+                 hooks,
+                 nf=512,
+                 blur=True,
+                 last_norm='Spectral',
+                 num_queries=100,
+                 num_scales=3,
+                 dec_layers=9):
+        super().__init__()
+        self.hooks = hooks
+        self.nf = nf
+        self.blur = blur
+        self.last_norm = getattr(NormType, last_norm)
+
+        self.layers = self.make_layers()
+        embed_dim = nf // 2
+
+        self.last_shuf = CustomPixelShuffle_ICNR(
+            embed_dim,
+            embed_dim,
+            blur=self.blur,
+            norm_type=self.last_norm,
+            scale=4)
+
+        self.color_decoder = MultiScaleColorDecoder(
+            in_channels=[512, 512, 256],
+            num_queries=num_queries,
+            num_scales=num_scales,
+            dec_layers=dec_layers,
+        )
+
+    def forward(self):
+        encode_feat = self.hooks[-1].feature
+        out0 = self.layers[0](encode_feat)
+        out1 = self.layers[1](out0)
+        out2 = self.layers[2](out1)
+        out3 = self.last_shuf(out2)
+        out = self.color_decoder([out0, out1, out2], out3)
+
+        return out
+
+    def make_layers(self):
+        decoder_layers = []
+
+        e_in_c = self.hooks[-1].feature.shape[1]
+        in_c = e_in_c
+
+        out_c = self.nf
+        setup_hooks = self.hooks[-2::-1]
+        for layer_index, hook in enumerate(setup_hooks):
+            feature_c = hook.feature.shape[1]
+            if layer_index == len(setup_hooks) - 1:
+                out_c = out_c // 2
+            decoder_layers.append(
+                UnetBlockWide(
+                    in_c,
+                    feature_c,
+                    out_c,
+                    hook,
+                    blur=self.blur,
+                    self_attention=False,
+                    norm_type=NormType.Spectral))
+            in_c = out_c
+        return nn.Sequential(*decoder_layers)
+
+
+class Encoder(nn.Module):
+
+    def __init__(self, encoder_name, hook_names, **kwargs):
+        super().__init__()
+        if encoder_name == 'convnext-t' or encoder_name == 'convnext':
+            self.arch = ConvNeXt()
+        elif encoder_name == 'convnext-s':
+            self.arch = ConvNeXt(
+                depths=[3, 3, 27, 3], dims=[96, 192, 384, 768])
+        elif encoder_name == 'convnext-b':
+            self.arch = ConvNeXt(
+                depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024])
+        elif encoder_name == 'convnext-l':
+            self.arch = ConvNeXt(
+                depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536])
+        else:
+            raise NotImplementedError
+
+        self.hook_names = hook_names
+        self.hooks = self.setup_hooks()
+
+    def setup_hooks(self):
+        hooks = [Hook(self.arch._modules[name]) for name in self.hook_names]
+        return hooks
+
+    def forward(self, img):
+        return self.arch(img)
+
+
+class MultiScaleColorDecoder(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 hidden_dim=256,
+                 num_queries=100,
+                 nheads=8,
+                 dim_feedforward=2048,
+                 dec_layers=9,
+                 pre_norm=False,
+                 color_embed_dim=256,
+                 enforce_input_project=True,
+                 num_scales=3):
+        super().__init__()
+
+        # positional encoding
+        N_steps = hidden_dim // 2
+        self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
+
+        # define Transformer decoder
+        self.num_heads = nheads
+        self.num_layers = dec_layers
+        self.transformer_self_attention_layers = nn.ModuleList()
+        self.transformer_cross_attention_layers = nn.ModuleList()
+        self.transformer_ffn_layers = nn.ModuleList()
+
+        for _ in range(self.num_layers):
+            self.transformer_self_attention_layers.append(
+                SelfAttentionLayer(
+                    d_model=hidden_dim,
+                    nhead=nheads,
+                    dropout=0.0,
+                    normalize_before=pre_norm,
+                ))
+            self.transformer_cross_attention_layers.append(
+                CrossAttentionLayer(
+                    d_model=hidden_dim,
+                    nhead=nheads,
+                    dropout=0.0,
+                    normalize_before=pre_norm,
+                ))
+            self.transformer_ffn_layers.append(
+                FFNLayer(
+                    d_model=hidden_dim,
+                    dim_feedforward=dim_feedforward,
+                    dropout=0.0,
+                    normalize_before=pre_norm,
+                ))
+
+        self.decoder_norm = nn.LayerNorm(hidden_dim)
+
+        self.num_queries = num_queries
+        # learnable color query features
+        self.query_feat = nn.Embedding(num_queries, hidden_dim)
+        # learnable color query p.e.
+        self.query_embed = nn.Embedding(num_queries, hidden_dim)
+
+        # level embedding
+        self.num_feature_levels = num_scales
+        self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim)
+
+        # input projections
+        self.input_proj = nn.ModuleList()
+        for i in range(self.num_feature_levels):
+            if in_channels[i] != hidden_dim or enforce_input_project:
+                self.input_proj.append(
+                    nn.Conv2d(in_channels[i], hidden_dim, kernel_size=1))
+                nn.init.kaiming_uniform_(self.input_proj[-1].weight, a=1)
+                if self.input_proj[-1].bias is not None:
+                    nn.init.constant_(self.input_proj[-1].bias, 0)
+            else:
+                self.input_proj.append(nn.Sequential())
+
+        # output FFNs
+        self.color_embed = MLP(hidden_dim, hidden_dim, color_embed_dim, 3)
+
+    def forward(self, feature_pyramid, last_img_feature):
+        assert len(feature_pyramid) == self.num_feature_levels
+        src, pos = [], []
+
+        for i in range(self.num_feature_levels):
+            pos.append(self.pe_layer(feature_pyramid[i], None).flatten(2))
+            src.append(self.input_proj[i](feature_pyramid[i]).flatten(2)
+                       + self.level_embed.weight[i][None, :, None])
+
+            # flatten NxCxHxW to HWxNxC
+            pos[-1] = pos[-1].permute(2, 0, 1)
+            src[-1] = src[-1].permute(2, 0, 1)
+
+        _, bs, _ = src[0].shape
+
+        # QxNxC
+        query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
+        output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)
+
+        for i in range(self.num_layers):
+            level_index = i % self.num_feature_levels
+            # attention: cross-attention first
+            output = self.transformer_cross_attention_layers[i](
+                output,
+                src[level_index],
+                memory_mask=None,
+                memory_key_padding_mask=None,
+                pos=pos[level_index],
+                query_pos=query_embed)
+            output = self.transformer_self_attention_layers[i](
+                output,
+                tgt_mask=None,
+                tgt_key_padding_mask=None,
+                query_pos=query_embed)
+            # FFN
+            output = self.transformer_ffn_layers[i](output)
+
+        decoder_output = self.decoder_norm(output)
+        decoder_output = decoder_output.transpose(
+            0, 1)  # [N, bs, C]  -> [bs, N, C]
+        color_embed = self.color_embed(decoder_output)
+        out = torch.einsum('bqc,bchw->bqhw', color_embed, last_img_feature)
+
+        return out

modelscope/models/cv/image_colorization/ddcolor/ddcolor_for_image_colorization.py

@@ -0,0 +1,52 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os.path as osp
+from typing import Dict, Union
+
+from modelscope.metainfo import Models
+from modelscope.models.base import Tensor, TorchModel
+from modelscope.models.builder import MODELS
+from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.logger import get_logger
+from .ddcolor import DDColor
+
+logger = get_logger()
+
+__all__ = ['DDColorForImageColorization']
+
+
+@MODELS.register_module(Tasks.image_colorization, module_name=Models.ddcolor)
+class DDColorForImageColorization(TorchModel):
+
+    def __init__(self,
+                 model_dir,
+                 encoder_name='convnext-l',
+                 input_size=(512, 512),
+                 num_queries=100,
+                 *args,
+                 **kwargs):
+        """initialize the image colorization model from the `model_dir` path.
+
+        Args:
+            model_dir (str): the model path.
+            encoder_name (str): the encoder name.
+            input_size (tuple): size of the model input image.
+            num_queries (int): number of decoder queries
+        """
+        super().__init__(model_dir, *args, **kwargs)
+
+        self.model = DDColor(encoder_name, input_size, num_queries)
+
+        model_path = osp.join(model_dir, ModelFile.TORCH_MODEL_FILE)
+        self.model = self._load_pretrained(self.model, model_path)
+
+    def forward(self, input: Dict[str,
+                                  Tensor]) -> Dict[str, Union[list, Tensor]]:
+        """return the result of the model
+
+        Args:
+            inputs (Tensor): the preprocessed data
+
+        Returns:
+            Dict[str, Tensor]: results
+        """
+        return self.model(**input)

modelscope/models/cv/image_colorization/ddcolor/utils/convnext.py

@@ -0,0 +1,177 @@
+# The implementation here is modified based on ConvNeXt, originally MIT license
+# and publicly available at https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+
+
+class Block(nn.Module):
+    r""" ConvNeXt Block. There are two equivalent implementations:
+    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+    We use (2) as we find it slightly faster in PyTorch
+
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+    """
+
+    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
+        super().__init__()
+        self.dwconv = nn.Conv2d(
+            dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(
+            dim,
+            4 * dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = nn.Parameter(
+            layer_scale_init_value * torch.ones((dim)),
+            requires_grad=True) if layer_scale_init_value > 0 else None
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+
+class ConvNeXt(nn.Module):
+    r""" ConvNeXt
+        A PyTorch impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+    Args:
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(
+        self,
+        in_chans=3,
+        num_classes=1000,
+        depths=[3, 3, 9, 3],
+        dims=[96, 192, 384, 768],
+        drop_path_rate=0.,
+        layer_scale_init_value=1e-6,
+        head_init_scale=1.,
+    ):
+        super().__init__()
+
+        self.downsample_layers = nn.ModuleList(
+        )  # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+            LayerNorm(dims[0], eps=1e-6, data_format='channels_first'))
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                LayerNorm(dims[i], eps=1e-6, data_format='channels_first'),
+                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = nn.ModuleList(
+        )  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [
+            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
+        ]
+        cur = 0
+        for i in range(4):
+            stage = nn.Sequential(*[
+                Block(
+                    dim=dims[i],
+                    drop_path=dp_rates[cur + j],
+                    layer_scale_init_value=layer_scale_init_value)
+                for j in range(depths[i])
+            ])
+            self.stages.append(stage)
+            cur += depths[i]
+
+        # add norm layers for each output
+        out_indices = (0, 1, 2, 3)
+        for i in out_indices:
+            layer = LayerNorm(dims[i], eps=1e-6, data_format='channels_first')
+            # layer = nn.Identity()
+            layer_name = f'norm{i}'
+            self.add_module(layer_name, layer)
+
+        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer
+        self.head_cls = nn.Linear(dims[-1], 4)
+
+        self.apply(self._init_weights)
+        self.head_cls.weight.data.mul_(head_init_scale)
+        self.head_cls.bias.data.mul_(head_init_scale)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward_features(self, x):
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+
+            # add extra norm
+            norm_layer = getattr(self, f'norm{i}')
+            norm_layer(x)
+
+        return self.norm(x.mean(
+            [-2, -1]))  # global average pooling, (N, C, H, W) -> (N, C)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head_cls(x)
+        return x
+
+
+class LayerNorm(nn.Module):
+    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+    with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self,
+                 normalized_shape,
+                 eps=1e-6,
+                 data_format='channels_last'):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ['channels_last', 'channels_first']:
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape, )
+
+    def forward(self, x):
+        if self.data_format == 'channels_last':  # B H W C
+            return F.layer_norm(x, self.normalized_shape, self.weight,
+                                self.bias, self.eps)
+        elif self.data_format == 'channels_first':  # B C H W
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x

modelscope/models/cv/image_colorization/ddcolor/utils/position_encoding.py

@@ -0,0 +1,57 @@
+# The implementation here is modified based on Mask2Former, originally MIT license and publicly available at
+# https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/position_encoding.py
+
+import math
+
+import torch
+from torch import nn
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+
+    def __init__(self,
+                 num_pos_feats=64,
+                 temperature=10000,
+                 normalize=False,
+                 scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError('normalize should be True if scale is passed')
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, x, mask=None):
+        if mask is None:
+            mask = torch.zeros((x.size(0), x.size(2), x.size(3)),
+                               device=x.device,
+                               dtype=torch.bool)
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(
+            self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature**(2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack(
+            (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()),
+            dim=4).flatten(3)
+        pos_y = torch.stack(
+            (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()),
+            dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        return pos

modelscope/models/cv/image_colorization/ddcolor/utils/transformer_utils.py

@@ -0,0 +1,232 @@
+# The implementation here is modified based on Mask2Former, originally MIT license and publicly available at
+# https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/mask2former_transformer_decoder.py
+
+from typing import Optional
+
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+class SelfAttentionLayer(nn.Module):
+
+    def __init__(self,
+                 d_model,
+                 nhead,
+                 dropout=0.0,
+                 activation='relu',
+                 normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+
+        self.norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self,
+                     tgt,
+                     tgt_mask: Optional[Tensor] = None,
+                     tgt_key_padding_mask: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None):
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2 = self.self_attn(
+            q,
+            k,
+            value=tgt,
+            attn_mask=tgt_mask,
+            key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout(tgt2)
+        tgt = self.norm(tgt)
+
+        return tgt
+
+    def forward_pre(self,
+                    tgt,
+                    tgt_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(
+            q,
+            k,
+            value=tgt2,
+            attn_mask=tgt_mask,
+            key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout(tgt2)
+
+        return tgt
+
+    def forward(self,
+                tgt,
+                tgt_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, tgt_mask, tgt_key_padding_mask,
+                                    query_pos)
+        return self.forward_post(tgt, tgt_mask, tgt_key_padding_mask,
+                                 query_pos)
+
+
+class CrossAttentionLayer(nn.Module):
+
+    def __init__(self,
+                 d_model,
+                 nhead,
+                 dropout=0.0,
+                 activation='relu',
+                 normalize_before=False):
+        super().__init__()
+        self.multihead_attn = nn.MultiheadAttention(
+            d_model, nhead, dropout=dropout)
+
+        self.norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self,
+                     tgt,
+                     memory,
+                     memory_mask: Optional[Tensor] = None,
+                     memory_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None):
+        tgt2 = self.multihead_attn(
+            query=self.with_pos_embed(tgt, query_pos),
+            key=self.with_pos_embed(memory, pos),
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout(tgt2)
+        tgt = self.norm(tgt)
+
+        return tgt
+
+    def forward_pre(self,
+                    tgt,
+                    memory,
+                    memory_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm(tgt)
+        tgt2 = self.multihead_attn(
+            query=self.with_pos_embed(tgt2, query_pos),
+            key=self.with_pos_embed(memory, pos),
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout(tgt2)
+
+        return tgt
+
+    def forward(self,
+                tgt,
+                memory,
+                memory_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory, memory_mask,
+                                    memory_key_padding_mask, pos, query_pos)
+        return self.forward_post(tgt, memory, memory_mask,
+                                 memory_key_padding_mask, pos, query_pos)
+
+
+class FFNLayer(nn.Module):
+
+    def __init__(self,
+                 d_model,
+                 dim_feedforward=2048,
+                 dropout=0.0,
+                 activation='relu',
+                 normalize_before=False):
+        super().__init__()
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm = nn.LayerNorm(d_model)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt):
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout(tgt2)
+        tgt = self.norm(tgt)
+        return tgt
+
+    def forward_pre(self, tgt):
+        tgt2 = self.norm(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout(tgt2)
+        return tgt
+
+    def forward(self, tgt):
+        if self.normalize_before:
+            return self.forward_pre(tgt)
+        return self.forward_post(tgt)
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == 'relu':
+        return F.relu
+    if activation == 'gelu':
+        return F.gelu
+    if activation == 'glu':
+        return F.glu
+    raise RuntimeError(F'activation should be relu/gelu, not {activation}.')
+
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x

modelscope/models/cv/image_colorization/ddcolor/utils/unet.py

@@ -0,0 +1,203 @@
+# The implementation here is modified based on DeOldify, originally MIT License
+# and publicly available at https://github.com/jantic/DeOldify/blob/master/deoldify/unet.py
+
+import collections
+from enum import Enum
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+NormType = Enum('NormType', 'Batch BatchZero Weight Spectral')
+
+
+class Hook:
+    feature = None
+
+    def __init__(self, module):
+        self.hook = module.register_forward_hook(self.hook_fn)
+
+    def hook_fn(self, module, input, output):
+        if isinstance(output, torch.Tensor):
+            self.feature = output
+        elif isinstance(output, collections.OrderedDict):
+            self.feature = output['out']
+
+    def remove(self):
+        self.hook.remove()
+
+
+class SelfAttention(nn.Module):
+    'Self attention layer for nd.'
+
+    def __init__(self, n_channels: int):
+        super().__init__()
+        self.query = conv1d(n_channels, n_channels // 8)
+        self.key = conv1d(n_channels, n_channels // 8)
+        self.value = conv1d(n_channels, n_channels)
+        self.gamma = nn.Parameter(torch.tensor([0.]))
+
+    def forward(self, x):
+        # Notation from https://arxiv.org/pdf/1805.08318.pdf
+        size = x.size()
+        x = x.view(*size[:2], -1)
+        f, g, h = self.query(x), self.key(x), self.value(x)
+        beta = F.softmax(torch.bmm(f.permute(0, 2, 1).contiguous(), g), dim=1)
+        o = self.gamma * torch.bmm(h, beta) + x
+        return o.view(*size).contiguous()
+
+
+def batchnorm_2d(nf: int, norm_type: NormType = NormType.Batch):
+    'A batchnorm2d layer with `nf` features initialized depending on `norm_type`.'
+    bn = nn.BatchNorm2d(nf)
+    with torch.no_grad():
+        bn.bias.fill_(1e-3)
+        bn.weight.fill_(0. if norm_type == NormType.BatchZero else 1.)
+    return bn
+
+
+def init_default(m: nn.Module, func=nn.init.kaiming_normal_) -> None:
+    'Initialize `m` weights with `func` and set `bias` to 0.'
+    if func:
+        if hasattr(m, 'weight'):
+            func(m.weight)
+        if hasattr(m, 'bias') and hasattr(m.bias, 'data'):
+            m.bias.data.fill_(0.)
+    return m
+
+
+def icnr(x, scale=2, init=nn.init.kaiming_normal_):
+    'ICNR init of `x`, with `scale` and `init` function.'
+    ni, nf, h, w = x.shape
+    ni2 = int(ni / (scale**2))
+    k = init(torch.zeros([ni2, nf, h, w])).transpose(0, 1)
+    k = k.contiguous().view(ni2, nf, -1)
+    k = k.repeat(1, 1, scale**2)
+    k = k.contiguous().view([nf, ni, h, w]).transpose(0, 1)
+    x.data.copy_(k)
+
+
+def conv1d(ni: int,
+           no: int,
+           ks: int = 1,
+           stride: int = 1,
+           padding: int = 0,
+           bias: bool = False):
+    'Create and initialize a `nn.Conv1d` layer with spectral normalization.'
+    conv = nn.Conv1d(ni, no, ks, stride=stride, padding=padding, bias=bias)
+    nn.init.kaiming_normal_(conv.weight)
+    if bias:
+        conv.bias.data.zero_()
+    return nn.utils.spectral_norm(conv)
+
+
+def custom_conv_layer(
+    ni: int,
+    nf: int,
+    ks: int = 3,
+    stride: int = 1,
+    padding: int = None,
+    bias: bool = None,
+    is_1d: bool = False,
+    norm_type=NormType.Batch,
+    use_activ: bool = True,
+    transpose: bool = False,
+    init=nn.init.kaiming_normal_,
+    self_attention: bool = False,
+    extra_bn: bool = False,
+):
+    'Create a sequence of convolutional (`ni` to `nf`), ReLU (if `use_activ`) and batchnorm (if `bn`) layers.'
+    if padding is None:
+        padding = (ks - 1) // 2 if not transpose else 0
+    bn = norm_type in (NormType.Batch, NormType.BatchZero) or extra_bn
+    if bias is None:
+        bias = not bn
+    conv_func = nn.ConvTranspose2d if transpose else nn.Conv1d if is_1d else nn.Conv2d
+    conv = init_default(
+        conv_func(
+            ni, nf, kernel_size=ks, bias=bias, stride=stride, padding=padding),
+        init,
+    )
+
+    if norm_type == NormType.Weight:
+        conv = nn.utils.weight_norm(conv)
+    elif norm_type == NormType.Spectral:
+        conv = nn.utils.spectral_norm(conv)
+    layers = [conv]
+    if use_activ:
+        layers.append(nn.ReLU(True))
+    if bn:
+        layers.append((nn.BatchNorm1d if is_1d else nn.BatchNorm2d)(nf))
+    if self_attention:
+        layers.append(SelfAttention(nf))
+    return nn.Sequential(*layers)
+
+
+class CustomPixelShuffle_ICNR(nn.Module):
+    """
+    Upsample by `scale` from `ni` filters to `nf` (default `ni`),
+    using `nn.PixelShuffle`, `icnr` init, and `weight_norm`.
+    """
+
+    def __init__(self,
+                 ni: int,
+                 nf: int = None,
+                 scale: int = 2,
+                 blur: bool = True,
+                 norm_type=NormType.Spectral,
+                 extra_bn=False):
+        super().__init__()
+        self.conv = custom_conv_layer(
+            ni,
+            nf * (scale**2),
+            ks=1,
+            use_activ=False,
+            norm_type=norm_type,
+            extra_bn=extra_bn)
+        icnr(self.conv[0].weight)
+        self.shuf = nn.PixelShuffle(scale)
+        self.do_blur = blur
+        # Blurring over (h*w) kernel
+        # "Super-Resolution using Convolutional Neural Networks without Any Checkerboard Artifacts"
+        # - https://arxiv.org/abs/1806.02658
+        self.pad = nn.ReplicationPad2d((1, 0, 1, 0))
+        self.blur = nn.AvgPool2d(2, stride=1)
+        self.relu = nn.ReLU(True)
+
+    def forward(self, x):
+        x = self.shuf(self.relu(self.conv(x)))
+        return self.blur(self.pad(x)) if self.do_blur else x
+
+
+class UnetBlockWide(nn.Module):
+    'A quasi-UNet block, using `PixelShuffle_ICNR upsampling`.'
+
+    def __init__(self,
+                 up_in_c: int,
+                 x_in_c: int,
+                 n_out: int,
+                 hook,
+                 blur: bool = False,
+                 self_attention: bool = False,
+                 norm_type=NormType.Spectral):
+        super().__init__()
+
+        self.hook = hook
+        up_out = n_out
+        self.shuf = CustomPixelShuffle_ICNR(
+            up_in_c, up_out, blur=blur, norm_type=norm_type, extra_bn=True)
+        self.bn = batchnorm_2d(x_in_c)
+        ni = up_out + x_in_c
+        self.conv = custom_conv_layer(
+            ni,
+            n_out,
+            norm_type=norm_type,
+            self_attention=self_attention,
+            extra_bn=True)
+        self.relu = nn.ReLU()
+
+    def forward(self, up_in):
+        s = self.hook.feature
+        up_out = self.shuf(up_in)
+        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))
+        return self.conv(cat_x)

modelscope/models/cv/image_colorization/unet/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .unet import DynamicUnetDeep, DynamicUnetWide
+from .utils import NormType

modelscope/pipelines/builder.py

@@ -163,8 +163,8 @@ DEFAULT_MODEL_FOR_PIPELINE = {
      'damo/cv_csrnet_image-color-enhance-models'),
     Tasks.virtual_try_on: (Pipelines.virtual_try_on,
                            'damo/cv_daflow_virtual-try-on_base'),
-    Tasks.image_colorization: (Pipelines.image_colorization,
-                               'damo/cv_unet_image-colorization'),
+    Tasks.image_colorization: (Pipelines.ddcolor_image_colorization,
+                               'damo/cv_ddcolor_image-colorization'),
     Tasks.image_segmentation:
     (Pipelines.image_instance_segmentation,
      'damo/cv_swin-b_image-instance-segmentation_coco'),

modelscope/pipelines/cv/__init__.py

@@ -74,6 +74,7 @@ if TYPE_CHECKING:
     from .pointcloud_sceneflow_estimation_pipeline import PointCloudSceneFlowEstimationPipeline
     from .maskdino_instance_segmentation_pipeline import MaskDINOInstanceSegmentationPipeline
     from .image_mvs_depth_estimation_pipeline import ImageMultiViewDepthEstimationPipeline
+    from .ddcolor_image_colorization_pipeline import DDColorImageColorizationPipeline
 
 else:
     _import_structure = {
@@ -176,6 +177,9 @@ else:
         'image_mvs_depth_estimation_pipeline': [
             'ImageMultiViewDepthEstimationPipeline'
         ],
+        'ddcolor_image_colorization_pipeline': [
+            'DDColorImageColorizationPipeline'
+        ],
     }
 
     import sys

modelscope/pipelines/cv/ddcolor_image_colorization_pipeline.py

@@ -0,0 +1,168 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Any, Dict, Optional, Union
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+
+from modelscope.metainfo import Pipelines
+from modelscope.models.cv.image_colorization import DDColorForImageColorization
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines.base import Input, Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.preprocessors import LoadImage
+from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+
+@PIPELINES.register_module(
+    Tasks.image_colorization, module_name=Pipelines.ddcolor_image_colorization)
+class DDColorImageColorizationPipeline(Pipeline):
+    """ DDColor Image Colorization Pipeline.
+
+    Example:
+
+    ```python
+    >>> from modelscope.pipelines import pipeline
+
+    >>> colorizer = pipeline('image-colorization', 'damo/cv_ddcolor_image-colorization')
+    >>> colorizer("data/test/images/audrey_hepburn.jpg")
+       {'output_img': array([[[198, 199, 193],
+         [198, 199, 193],
+         [197, 199, 195],
+         ...,
+         [197, 213, 206],
+         [197, 213, 206],
+         [197, 213, 207]],
+
+        [[198, 199, 193],
+         [198, 199, 193],
+         [197, 199, 195],
+         ...,
+         [196, 212, 205],
+         [196, 212, 205],
+         [196, 212, 206]],
+
+        [[198, 199, 193],
+         [198, 199, 193],
+         [197, 199, 195],
+         ...,
+         [193, 209, 202],
+         [193, 209, 202],
+         [193, 209, 203]],
+
+        ...,
+
+        [[ 56,  72, 103],
+         [ 56,  72, 103],
+         [ 56,  72, 102],
+         ...,
+         [233, 231, 232],
+         [233, 231, 232],
+         [233, 231, 232]],
+
+        [[ 51,  62,  91],
+         [ 52,  63,  92],
+         [ 52,  64,  92],
+         ...,
+         [233, 232, 231],
+         [233, 232, 231],
+         [232, 232, 229]],
+
+        [[ 60,  72, 101],
+         [ 59,  71, 100],
+         [ 57,  70,  99],
+         ...,
+         [233, 232, 231],
+         [233, 232, 231],
+         [232, 232, 229]]], dtype=uint8)}
+    >>> #
+    ```
+    """
+
+    def __init__(self, model: Union[DDColorForImageColorization, str],
+                 **kwargs):
+        """
+        use `model` to create an image colorization pipeline for prediction
+
+        Args:
+            model: model id on modelscope hub.
+        """
+        super().__init__(model=model, **kwargs)
+        self.model.eval()
+        self.input_size = 512
+        if torch.cuda.is_available():
+            self._device = torch.device('cuda')
+        else:
+            self._device = torch.device('cpu')
+
+        # self.model = DDColorForImageColorization(
+        #     model_dir=model,
+        #     encoder_name='convnext-l',
+        #     input_size=[self.input_size, self.input_size],
+        # ).to(self.device)
+
+        # model_path = f'{model}/{ModelFile.TORCH_MODEL_FILE}'
+        # logger.info(f'loading model from {model_path}')
+        # self.model.load_state_dict(
+        #     torch.load(model_path, map_location=torch.device('cpu'))['params'],
+        #     strict=True)
+
+        logger.info('load model done')
+
+    def preprocess(self, input: Input) -> Dict[str, Any]:
+        """preprocess the input image, extract L-channel and convert it back to RGB
+
+        Args:
+            inputs: an input image from file or url
+
+        Returns:
+            Dict[str, Any]: the pre-processed image
+        """
+        img = LoadImage.convert_to_ndarray(input)
+        self.height, self.width = img.shape[:2]
+
+        img = (img / 255.0).astype(np.float32)
+        self.orig_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]
+
+        img = cv2.resize(img, (self.input_size, self.input_size))
+        img_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]
+        img_gray_lab = np.concatenate(
+            (img_l, np.zeros_like(img_l), np.zeros_like(img_l)), axis=-1)
+        img_gray_rgb = cv2.cvtColor(img_gray_lab, cv2.COLOR_LAB2RGB)
+        tensor_gray_rgb = torch.from_numpy(img_gray_rgb.transpose(
+            (2, 0, 1))).float()
+        tensor_gray_rgb = tensor_gray_rgb.unsqueeze(0).to(self.device)
+
+        result = {'img': tensor_gray_rgb}
+        return result
+
+    @torch.no_grad()
+    def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
+        """call model to output the predictions and concatenate it with the original L-channel
+
+        Args:
+            inputs: input image tensor
+
+        Returns:
+            Dict[str, Any]: the result image
+        """
+
+        output_ab = self.model(input).cpu()
+
+        output_ab_resize = F.interpolate(
+            output_ab, size=(self.height, self.width))
+        output_ab_resize = output_ab_resize[0].float().numpy().transpose(
+            1, 2, 0)
+        out_lab = np.concatenate((self.orig_l, output_ab_resize), axis=-1)
+        out_bgr = cv2.cvtColor(out_lab, cv2.COLOR_LAB2BGR)
+        output_img = (out_bgr * 255.0).round().astype(np.uint8)
+
+        return {OutputKeys.OUTPUT_IMG: output_img}
+
+    def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        return inputs

tests/pipelines/test_ddcolor_image_colorization.py

@@ -0,0 +1,61 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os.path as osp
+import unittest
+
+import cv2
+
+from modelscope.hub.snapshot_download import snapshot_download
+from modelscope.models import Model
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.pipelines.base import Pipeline
+from modelscope.pipelines.cv import DDColorImageColorizationPipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.demo_utils import DemoCompatibilityCheck
+from modelscope.utils.test_utils import test_level
+
+
+class DDColorImageColorizationTest(unittest.TestCase, DemoCompatibilityCheck):
+
+    def setUp(self) -> None:
+        self.task = Tasks.image_colorization
+        self.model_id = 'damo/cv_ddcolor_image-colorization'
+        self.test_image = 'data/test/images/audrey_hepburn.jpg'
+
+    def pipeline_inference(self, pipeline: Pipeline, test_image: str):
+        result = pipeline(test_image)
+        if result is not None:
+            cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])
+            print(f'Output written to {osp.abspath("result.png")}')
+
+    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+    def test_run_by_direct_model_download(self):
+        cache_path = snapshot_download(self.model_id)
+        image_colorization = DDColorImageColorizationPipeline(cache_path)
+        self.pipeline_inference(image_colorization, self.test_image)
+
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+    def test_run_with_model_from_pretrained(self):
+        model = Model.from_pretrained(self.model_id)
+        image_colorization = pipeline(
+            task=Tasks.image_colorization, model=model)
+        self.pipeline_inference(image_colorization, self.test_image)
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_model_from_modelhub(self):
+        image_colorization = pipeline(
+            task=Tasks.image_colorization, model=self.model_id)
+        self.pipeline_inference(image_colorization, self.test_image)
+
+    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+    def test_run_with_default_model(self):
+        image_colorization = pipeline(Tasks.image_colorization)
+        self.pipeline_inference(image_colorization, self.test_image)
+
+    @unittest.skip('demo compatibility test is only enabled on a needed-basis')
+    def test_demo_compatibility(self):
+        self.compatibility_check()
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/run_config.yaml

@@ -45,6 +45,7 @@ isolated:  # test cases that may require excessive anmount of GPU memory or run
   - test_video_super_resolution.py
   - test_kws_nearfield_trainer.py
   - test_gpt3_text_generation.py
+  - test_ddcolor_image_colorization.py
 
 envs:
   default: # default env, case not in other env will in default, pytorch.