fix importing issues

models/SRFlow/35000_G.pth

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

models/SRFlow/code/Measure.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import os
+import time
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import cv2
+import argparse
+
+from natsort import natsort
+from skimage.metrics import structural_similarity as ssim
+from skimage.metrics import peak_signal_noise_ratio as psnr
+import lpips
+
+
+class Measure():
+    def __init__(self, net='alex', use_gpu=False):
+        self.device = 'cuda' if use_gpu else 'cpu'
+        self.model = lpips.LPIPS(net=net)
+        self.model.to(self.device)
+
+    def measure(self, imgA, imgB):
+        return [float(f(imgA, imgB)) for f in [self.psnr, self.ssim, self.lpips]]
+
+    def lpips(self, imgA, imgB, model=None):
+        tA = t(imgA).to(self.device)
+        tB = t(imgB).to(self.device)
+        dist01 = self.model.forward(tA, tB).item()
+        return dist01
+
+    def ssim(self, imgA, imgB):
+        # multichannel: If True, treat the last dimension of the array as channels. Similarity calculations are done independently for each channel then averaged.
+        score, diff = ssim(imgA, imgB, full=True, multichannel=True, channel_axis=-1)
+        return score
+
+    def psnr(self, imgA, imgB):
+        psnr_val = psnr(imgA, imgB)
+        return psnr_val
+
+
+def t(img):
+    def to_4d(img):
+        assert len(img.shape) == 3
+        assert img.dtype == np.uint8
+        img_new = np.expand_dims(img, axis=0)
+        assert len(img_new.shape) == 4
+        return img_new
+
+    def to_CHW(img):
+        return np.transpose(img, [2, 0, 1])
+
+    def to_tensor(img):
+        return torch.Tensor(img)
+
+    return to_tensor(to_4d(to_CHW(img))) / 127.5 - 1
+
+
+def fiFindByWildcard(wildcard):
+    return natsort.natsorted(glob.glob(wildcard, recursive=True))
+
+
+def imread(path):
+    return cv2.imread(path)[:, :, [2, 1, 0]]
+
+
+def format_result(psnr, ssim, lpips):
+    return f'{psnr:0.2f}, {ssim:0.3f}, {lpips:0.3f}'
+
+def measure_dirs(dirA, dirB, use_gpu, verbose=False):
+    if verbose:
+        vprint = lambda x: print(x)
+    else:
+        vprint = lambda x: None
+
+
+    t_init = time.time()
+
+    paths_A = fiFindByWildcard(os.path.join(dirA, f'*.{type}'))
+    paths_B = fiFindByWildcard(os.path.join(dirB, f'*.{type}'))
+
+    vprint("Comparing: ")
+    vprint(dirA)
+    vprint(dirB)
+
+    measure = Measure(use_gpu=use_gpu)
+
+    results = []
+    for pathA, pathB in zip(paths_A, paths_B):
+        result = OrderedDict()
+
+        t = time.time()
+        result['psnr'], result['ssim'], result['lpips'] = measure.measure(imread(pathA), imread(pathB))
+        d = time.time() - t
+        vprint(f"{pathA.split('/')[-1]}, {pathB.split('/')[-1]}, {format_result(**result)}, {d:0.1f}")
+
+        results.append(result)
+
+    psnr = np.mean([result['psnr'] for result in results])
+    ssim = np.mean([result['ssim'] for result in results])
+    lpips = np.mean([result['lpips'] for result in results])
+
+    vprint(f"Final Result: {format_result(psnr, ssim, lpips)}, {time.time() - t_init:0.1f}s")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-dirA', default='', type=str)
+    parser.add_argument('-dirB', default='', type=str)
+    parser.add_argument('-type', default='png')
+    parser.add_argument('--use_gpu', action='store_true', default=False)
+    args = parser.parse_args()
+
+    dirA = args.dirA
+    dirB = args.dirB
+    type = args.type
+    use_gpu = args.use_gpu
+
+    if len(dirA) > 0 and len(dirB) > 0:
+        measure_dirs(dirA, dirB, use_gpu=use_gpu, verbose=True)

models/SRFlow/code/__init__.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+
+import glob 
+import sys
+
+sys.path.append('../..')
+from natsort import natsort
+import SRFlow.code.options.options as option
+
+from SRFlow.code.models import create_model
+import torch
+from SRFlow.code.utils.util import opt_get
+from SRFlow.code.models.SRFlow_model import SRFlowModel
+import numpy as np
+import os
+import cv2
+
+
+def fiFindByWildcard(wildcard):
+    return natsort.natsorted(glob.glob(wildcard, recursive=True))
+
+
+def load_model(conf_path):
+    opt = option.parse(conf_path, is_train=False)
+    opt['gpu_ids'] = None
+    opt = option.dict_to_nonedict(opt)
+    model = SRFlowModel(opt, 0)
+
+    model_path = opt_get(opt, ['model_path'], None)
+    model.load_network(load_path='models/SRFlow/35000_G.pth', network=model.netG)
+    return model, opt
+
+
+def predict(model, lr):
+    model.feed_data({"LQ": t(lr)}, need_GT=False)
+    model.test()
+    visuals = model.get_current_visuals(need_GT=False)
+    return visuals.get('rlt', visuals.get("SR"))
+
+
+def t(array): return torch.Tensor(np.expand_dims(array.transpose([2, 0, 1]), axis=0).astype(np.float32)) / 255
+
+
+def rgb(t): return (
+        np.clip((t[0] if len(t.shape) == 4 else t).detach().cpu().numpy().transpose([1, 2, 0]), 0, 1) * 255).astype(
+    np.uint8)
+
+
+def imread(path):
+    return cv2.imread(path)[:, :, [2, 1, 0]]
+
+
+def imwrite(path, img):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    cv2.imwrite(path, img[:, :, [2, 1, 0]])
+
+
+def imCropCenter(img, size):
+    h, w, c = img.shape
+
+    h_start = max(h // 2 - size // 2, 0)
+    h_end = min(h_start + size, h)
+
+    w_start = max(w // 2 - size // 2, 0)
+    w_end = min(w_start + size, w)
+
+    return img[h_start:h_end, w_start:w_end]
+
+
+def impad(img, top=0, bottom=0, left=0, right=0, color=255):
+    return np.pad(img, [(top, bottom), (left, right), (0, 0)], 'reflect')

models/SRFlow/code/a.py

@@ -0,0 +1,27 @@
+import pickle
+import numpy as np
+import os
+import matplotlib.pyplot as plt
+
+def load_pkls(path):
+    assert os.path.isfile(path), path
+    images = []
+    with open(path, "rb") as f:
+        images += pickle.load(f)
+    assert len(images) > 0, path
+    images = [np.transpose(image, [2, 0, 1]) for image in images]
+    return images
+
+path = 'datasets/DIV2K-va.pklv4'
+loaded_images = load_pkls(path)
+print(len(loaded_images))
+# Display the first image
+if loaded_images:
+    first_image = loaded_images[11]
+    plt.imshow(np.transpose(first_image, [1, 2, 0]))  # Transpose image to original shape [height, width, channels]
+    plt.title('First Image')
+    plt.axis('off')  # Hide axis
+    plt.show()
+else:
+    print("No images loaded from the pickle file.")
+print(loaded_images[11])

models/SRFlow/code/confs/RRDB_CelebA_8X.yml

@@ -0,0 +1,83 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SR
+distortion: sr
+scale: 8
+#gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/celebA-train-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/celebA-train-x8_1pct.pklv4
+
+    use_shuffle: true
+    n_workers: 0  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    use_rot: true
+    color: RGB
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/celebA-valid-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/celebA-valid-x8_1pct.pklv4
+    n_max: 10
+
+#### network structures
+network_G:
+  which_model_G: RRDBNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+
+#### path
+path:
+  pretrain_model_G: ~
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  lr_G: !!float 2e-4
+  lr_scheme: CosineAnnealingLR_Restart
+  beta1: 0.9
+  beta2: 0.99
+  niter: 200000
+  warmup_iter: -1  # no warm up
+  T_period: [ 50000, 50000, 50000, 50000 ]
+  restarts: [ 50000, 100000, 150000 ]
+  restart_weights: [ 1, 1, 1 ]
+  eta_min: !!float 1e-7
+
+  pixel_criterion: l1
+  pixel_weight: 1.0
+
+  manual_seed: 10
+  val_freq: !!float 5e3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/confs/RRDB_DF2K_4X.yml

@@ -0,0 +1,85 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SR
+distortion: sr
+scale: 4
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-train-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/DF2K-train-x4_1pct.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    color: RGB
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-valid-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/DF2K-valid-x4_1pct.pklv4
+    quant: 32
+    n_max: 20
+
+#### network structures
+network_G:
+  which_model_G: RRDBNet
+  use_orig: True
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+
+#### path
+path:
+  pretrain_model_G: ~
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  lr_G: !!float 2e-4
+  lr_scheme: CosineAnnealingLR_Restart
+  beta1: 0.9
+  beta2: 0.99
+  niter: 1000000
+  warmup_iter: -1  # no warm up
+  T_period: [ 50000, 50000, 50000, 50000 ]
+  restarts: [ 50000, 100000, 150000 ]
+  restart_weights: [ 1, 1, 1 ]
+  eta_min: !!float 1e-7
+
+  pixel_criterion: l1
+  pixel_weight: 1.0
+
+  manual_seed: 10
+  val_freq: !!float 5e3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/confs/RRDB_DF2K_8X.yml

@@ -0,0 +1,85 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SR
+distortion: sr
+scale: 8
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-train-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/DF2K-train-x8_1pct.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    color: RGB
+
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-valid-gt_1pct.pklv4
+    dataroot_LQ: ../datasets/DF2K-valid-x8_1pct.pklv4
+    quant: 32
+    n_max: 20
+
+#### network structures
+network_G:
+  which_model_G: RRDBNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+
+#### path
+path:
+  pretrain_model_G: ~
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  lr_G: !!float 2e-4
+  lr_scheme: CosineAnnealingLR_Restart
+  beta1: 0.9
+  beta2: 0.99
+  niter: 200000
+  warmup_iter: -1  # no warm up
+  T_period: [ 50000, 50000, 50000, 50000 ]
+  restarts: [ 50000, 100000, 150000 ]
+  restart_weights: [ 1, 1, 1 ]
+  eta_min: !!float 1e-7
+
+  pixel_criterion: l1
+  pixel_weight: 1.0
+
+  manual_seed: 10
+  val_freq: !!float 5e3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/confs/SRFlow_CelebA_8X.yml

@@ -0,0 +1,107 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SRFlow
+distortion: sr
+scale: 8
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/celebA-train-gt.pklv4
+    dataroot_LQ: ../datasets/celebA-train-x8.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    color: RGB
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/celebA-train-gt.pklv4
+    dataroot_LQ: ../datasets/celebA-train-x8.pklv4
+    quant: 32
+    n_max: 20
+
+#### Test Settings
+dataroot_GT: ../datasets/celebA-validation-gt
+dataroot_LR: ../datasets/celebA-validation-x8
+model_path: ../pretrained_models/SRFlow_CelebA_8X.pth
+heat: 0.9 # This is the standard deviation of the latent vectors
+
+#### network structures
+network_G:
+  which_model_G: SRFlowNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 8
+  upscale: 8
+  train_RRDB: false
+  train_RRDB_delay: 0.5
+
+  flow:
+    K: 16
+    L: 4
+    noInitialInj: true
+    coupling: CondAffineSeparatedAndCond
+    additionalFlowNoAffine: 2
+    split:
+      enable: true
+    fea_up0: true
+    stackRRDB:
+      blocks: [ 1, 3, 5, 7 ]
+      concat: true
+
+#### path
+path:
+  pretrain_model_G: ../pretrained_models/RRDB_CelebA_8X.pth
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  manual_seed: 10
+  lr_G: !!float 5e-4
+  weight_decay_G: 0
+  beta1: 0.9
+  beta2: 0.99
+  lr_scheme: MultiStepLR
+  warmup_iter: -1  # no warm up
+  lr_steps_rel: [ 0.5, 0.75, 0.9, 0.95 ]
+  lr_gamma: 0.5
+
+  niter: 200000
+  val_freq: 40000
+
+#### validation settings
+val:
+  heats: [ 0.0, 0.5, 0.75, 1.0 ]
+  n_sample: 3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/confs/SRFlow_DF2K_4X.yml

@@ -0,0 +1,106 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SRFlow
+distortion: sr
+scale: 4
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: DF2K_256_tr
+    mode: LRHR_PKL
+    dataroot_GT: /kaggle/input/srflow0103/SRFlow/datasets/DF2K-tr.pklv4
+    dataroot_LQ: /kaggle/input/srflow0103/SRFlow/datasets/DF2K-tr_X4.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 12
+    GT_size: 256
+    use_flip: true
+    color: RGB
+  val:
+    name: DF2K_256_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DIV2K-va.pklv4
+    dataroot_LQ: ../datasets/DIV2K-va_X4.pklv4
+    quant: 32
+    n_max: 20
+
+#### Test Settings
+dataroot: /kaggle/input/test-set/test set
+model_path: models/SRFlow/35000_G
+heat: 0.6 # This is the standard deviation of the latent vectors
+
+#### network structures
+network_G:
+  which_model_G: SRFlowNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+  upscale: 4
+  train_RRDB: false
+  train_RRDB_delay: 0.5
+
+  flow:
+    K: 16
+    L: 3
+    noInitialInj: true
+    coupling: CondAffineSeparatedAndCond
+    additionalFlowNoAffine: 2
+    split:
+      enable: true
+    fea_up0: true
+    stackRRDB:
+      blocks: [ 1, 8, 15, 22 ]
+      concat: true
+
+#### path
+path:
+  pretrain_model_G: 
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  manual_seed: 10
+  lr_G: !!float 2.5e-4
+  weight_decay_G: 0
+  beta1: 0.9
+  beta2: 0.99
+  lr_scheme: MultiStepLR
+  warmup_iter: -1  # no warm up
+  lr_steps_rel: [ 0.5, 0.75, 0.9, 0.95 ]
+  lr_gamma: 0.5
+
+  niter: 64185
+  val_freq: 40000
+
+#### validation settings
+val:
+  heats: [ 0.0, 0.5, 0.75, 1.0 ]
+  n_sample: 3
+
+#### logger
+logger:
+  print_freq: 100
+  save_checkpoint_freq: !!float 5e3

models/SRFlow/code/confs/SRFlow_DF2K_8X.yml

@@ -0,0 +1,112 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+#### general settings
+name: train
+use_tb_logger: true
+model: SRFlow
+distortion: sr
+scale: 8
+gpu_ids: [ 0 ]
+
+#### datasets
+datasets:
+  train:
+    name: CelebA_160_tr
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DF2K-tr.pklv4
+    dataroot_LQ: ../datasets/DF2K-tr_X8.pklv4
+    quant: 32
+
+    use_shuffle: true
+    n_workers: 3  # per GPU
+    batch_size: 16
+    GT_size: 160
+    use_flip: true
+    color: RGB
+
+  val:
+    name: CelebA_160_va
+    mode: LRHR_PKL
+    dataroot_GT: ../datasets/DIV2K-va.pklv4
+    dataroot_LQ: ../datasets/DIV2K-va_X8.pklv4
+    quant: 32
+    n_max: 20
+
+#### Test Settings
+dataroot_GT: ../datasets/div2k-validation-modcrop8-gt
+dataroot_LR: ../datasets/div2k-validation-modcrop8-x8
+model_path: ../pretrained_models/SRFlow_DF2K_8X.pth
+heat: 0.9 # This is the standard deviation of the latent vectors
+
+#### network structures
+network_G:
+  which_model_G: SRFlowNet
+  in_nc: 3
+  out_nc: 3
+  nf: 64
+  nb: 23
+  upscale: 8
+  train_RRDB: false
+  train_RRDB_delay: 0.5
+
+  flow:
+    K: 16
+    L: 4
+    noInitialInj: true
+    coupling: CondAffineSeparatedAndCond
+    additionalFlowNoAffine: 2
+    split:
+      enable: true
+    fea_up0: true
+    stackRRDB:
+      blocks: [ 1, 3, 5, 7 ]
+      concat: true
+
+#### path
+path:
+  pretrain_model_G: ../pretrained_models/RRDB_DF2K_8X.pth
+  strict_load: true
+  resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+  manual_seed: 10
+  lr_G: !!float 5e-4
+  weight_decay_G: 0
+  beta1: 0.9
+  beta2: 0.99
+  lr_scheme: MultiStepLR
+  warmup_iter: -1  # no warm up
+  lr_steps_rel: [ 0.5, 0.75, 0.9, 0.95 ]
+  lr_gamma: 0.5
+
+  niter: 200000
+  val_freq: 40000
+
+#### validation settings
+val:
+  heats: [ 0.0, 0.5, 0.75, 1.0 ]
+  n_sample: 3
+
+test:
+  heats: [ 0.0, 0.7, 0.8, 0.9 ]
+
+#### logger
+logger:
+  # Debug print_freq: 100
+  print_freq: 100
+  save_checkpoint_freq: !!float 1e3

models/SRFlow/code/data/LRHR_PKL_dataset.py

@@ -0,0 +1,179 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import os
+import subprocess
+import torch.utils.data as data
+import numpy as np
+import time
+import torch
+
+import pickle
+
+
+class LRHR_PKLDataset(data.Dataset):
+    def __init__(self, opt):
+        super(LRHR_PKLDataset, self).__init__()
+        self.opt = opt
+        self.crop_size = opt.get("GT_size", None)
+        self.scale = None
+        self.random_scale_list = [1]
+
+        hr_file_path = opt["dataroot_GT"]
+        lr_file_path = opt["dataroot_LQ"]
+        y_labels_file_path = opt['dataroot_y_labels']
+
+        gpu = True
+        augment = True
+
+        self.use_flip = opt["use_flip"] if "use_flip" in opt.keys() else False
+        self.use_rot = opt["use_rot"] if "use_rot" in opt.keys() else False
+        self.use_crop = opt["use_crop"] if "use_crop" in opt.keys() else False
+        self.center_crop_hr_size = opt.get("center_crop_hr_size", None)
+
+        n_max = opt["n_max"] if "n_max" in opt.keys() else int(1e8)
+
+        t = time.time()
+        self.lr_images = self.load_pkls(lr_file_path, n_max)
+        self.hr_images = self.load_pkls(hr_file_path, n_max)
+
+        min_val_hr = np.min([i.min() for i in self.hr_images[:20]])
+        max_val_hr = np.max([i.max() for i in self.hr_images[:20]])
+
+        min_val_lr = np.min([i.min() for i in self.lr_images[:20]])
+        max_val_lr = np.max([i.max() for i in self.lr_images[:20]])
+
+        t = time.time() - t
+        print("Loaded {} HR images with [{:.2f}, {:.2f}] in {:.2f}s from {}".
+              format(len(self.hr_images), min_val_hr, max_val_hr, t, hr_file_path))
+        print("Loaded {} LR images with [{:.2f}, {:.2f}] in {:.2f}s from {}".
+              format(len(self.lr_images), min_val_lr, max_val_lr, t, lr_file_path))
+
+        self.gpu = gpu
+        self.augment = augment
+
+        self.measures = None
+
+    def load_pkls(self, path, n_max):
+        assert os.path.isfile(path), path
+        images = []
+        with open(path, "rb") as f:
+            images += pickle.load(f)
+        assert len(images) > 0, path
+        images = images[:n_max]
+        images = [np.transpose(image, [2, 0, 1]) for image in images]
+        return images
+
+    def __len__(self):
+        return len(self.hr_images)
+
+    def __getitem__(self, item):
+        hr = self.hr_images[item]
+        lr = self.lr_images[item]
+
+        if self.scale == None:
+            self.scale = hr.shape[1] // lr.shape[1]
+            assert hr.shape[1] == self.scale * lr.shape[1], ('non-fractional ratio', lr.shape, hr.shape)
+
+        if self.use_crop:
+            hr, lr = random_crop(hr, lr, self.crop_size, self.scale, self.use_crop)
+
+        if self.center_crop_hr_size:
+            hr, lr = center_crop(hr, self.center_crop_hr_size), center_crop(lr, self.center_crop_hr_size // self.scale)
+
+        if self.use_flip:
+            hr, lr = random_flip(hr, lr)
+
+        if self.use_rot:
+            hr, lr = random_rotation(hr, lr)
+
+        hr = hr / 255.0
+        lr = lr / 255.0
+
+        if self.measures is None or np.random.random() < 0.05:
+            if self.measures is None:
+                self.measures = {}
+            self.measures['hr_means'] = np.mean(hr)
+            self.measures['hr_stds'] = np.std(hr)
+            self.measures['lr_means'] = np.mean(lr)
+            self.measures['lr_stds'] = np.std(lr)
+
+        hr = torch.Tensor(hr)
+        lr = torch.Tensor(lr)
+
+        # if self.gpu:
+        #    hr = hr.cuda()
+        #    lr = lr.cuda()
+
+        return {'LQ': lr, 'GT': hr, 'LQ_path': str(item), 'GT_path': str(item)}
+
+    def print_and_reset(self, tag):
+        m = self.measures
+        kvs = []
+        for k in sorted(m.keys()):
+            kvs.append("{}={:.2f}".format(k, m[k]))
+        print("[KPI] " + tag + ": " + ", ".join(kvs))
+        self.measures = None
+
+
+def random_flip(img, seg):
+    random_choice = np.random.choice([True, False])
+    img = img if random_choice else np.flip(img, 2).copy()
+    seg = seg if random_choice else np.flip(seg, 2).copy()
+    return img, seg
+
+
+def random_rotation(img, seg):
+    random_choice = np.random.choice([0, 1, 3])
+    img = np.rot90(img, random_choice, axes=(1, 2)).copy()
+    seg = np.rot90(seg, random_choice, axes=(1, 2)).copy()
+    return img, seg
+
+
+def random_crop(hr, lr, size_hr, scale, random):
+    size_lr = size_hr // scale
+
+    size_lr_x = lr.shape[1]
+    size_lr_y = lr.shape[2]
+
+    start_x_lr = np.random.randint(low=0, high=(size_lr_x - size_lr) + 1) if size_lr_x > size_lr else 0
+    start_y_lr = np.random.randint(low=0, high=(size_lr_y - size_lr) + 1) if size_lr_y > size_lr else 0
+
+    # LR Patch
+    lr_patch = lr[:, start_x_lr:start_x_lr + size_lr, start_y_lr:start_y_lr + size_lr]
+
+    # HR Patch
+    start_x_hr = start_x_lr * scale
+    start_y_hr = start_y_lr * scale
+    hr_patch = hr[:, start_x_hr:start_x_hr + size_hr, start_y_hr:start_y_hr + size_hr]
+
+    return hr_patch, lr_patch
+
+
+def center_crop(img, size):
+    assert img.shape[1] == img.shape[2], img.shape
+    border_double = img.shape[1] - size
+    assert border_double % 2 == 0, (img.shape, size)
+    border = border_double // 2
+    return img[:, border:-border, border:-border]
+
+
+def center_crop_tensor(img, size):
+    assert img.shape[2] == img.shape[3], img.shape
+    border_double = img.shape[2] - size
+    assert border_double % 2 == 0, (img.shape, size)
+    border = border_double // 2
+    return img[:, :, border:-border, border:-border]

models/SRFlow/code/data/__init__.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+'''create dataset and dataloader'''
+import logging
+import torch
+import torch.utils.data
+
+
+def create_dataloader(dataset, dataset_opt, opt=None, sampler=None):
+    phase = dataset_opt.get('phase', 'test')
+    if phase == 'train':
+        gpu_ids = opt.get('gpu_ids', None)
+        gpu_ids = gpu_ids if gpu_ids else []
+        num_workers = dataset_opt['n_workers'] * len(gpu_ids)
+        batch_size = dataset_opt['batch_size']
+        shuffle = True
+        return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
+                                           num_workers=num_workers, sampler=sampler, drop_last=True,
+                                           pin_memory=False)
+    else:
+        return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1,
+                                           pin_memory=True)
+
+
+def create_dataset(dataset_opt):
+    print(dataset_opt)
+    mode = dataset_opt['mode']
+    if mode == 'LRHR_PKL':
+        from data.LRHR_PKL_dataset import LRHR_PKLDataset as D
+    else:
+        raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode))
+    dataset = D(dataset_opt)
+
+    logger = logging.getLogger('base')
+    logger.info('Dataset [{:s} - {:s}] is created.'.format(dataset.__class__.__name__,
+                                                           dataset_opt['name']))
+    return dataset

models/SRFlow/code/imresize.py

@@ -0,0 +1,180 @@
+        # https://github.com/fatheral/matlab_imresize
+#
+# MIT License
+#
+# Copyright (c) 2020 Alex
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+from __future__ import print_function
+import numpy as np
+from math import ceil, floor
+
+
+def deriveSizeFromScale(img_shape, scale):
+    output_shape = []
+    for k in range(2):
+        output_shape.append(int(ceil(scale[k] * img_shape[k])))
+    return output_shape
+
+
+def deriveScaleFromSize(img_shape_in, img_shape_out):
+    scale = []
+    for k in range(2):
+        scale.append(1.0 * img_shape_out[k] / img_shape_in[k])
+    return scale
+
+
+def triangle(x):
+    x = np.array(x).astype(np.float64)
+    lessthanzero = np.logical_and((x >= -1), x < 0)
+    greaterthanzero = np.logical_and((x <= 1), x >= 0)
+    f = np.multiply((x + 1), lessthanzero) + np.multiply((1 - x), greaterthanzero)
+    return f
+
+
+def cubic(x):
+    x = np.array(x).astype(np.float64)
+    absx = np.absolute(x)
+    absx2 = np.multiply(absx, absx)
+    absx3 = np.multiply(absx2, absx)
+    f = np.multiply(1.5 * absx3 - 2.5 * absx2 + 1, absx <= 1) + np.multiply(-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2,
+                                                                            (1 < absx) & (absx <= 2))
+    return f
+
+
+def contributions(in_length, out_length, scale, kernel, k_width):
+    if scale < 1:
+        h = lambda x: scale * kernel(scale * x)
+        kernel_width = 1.0 * k_width / scale
+    else:
+        h = kernel
+        kernel_width = k_width
+    x = np.arange(1, out_length + 1).astype(np.float64)
+    u = x / scale + 0.5 * (1 - 1 / scale)
+    left = np.floor(u - kernel_width / 2)
+    P = int(ceil(kernel_width)) + 2
+    ind = np.expand_dims(left, axis=1) + np.arange(P) - 1  # -1 because indexing from 0
+    indices = ind.astype(np.int32)
+    weights = h(np.expand_dims(u, axis=1) - indices - 1)  # -1 because indexing from 0
+    weights = np.divide(weights, np.expand_dims(np.sum(weights, axis=1), axis=1))
+    aux = np.concatenate((np.arange(in_length), np.arange(in_length - 1, -1, step=-1))).astype(np.int32)
+    indices = aux[np.mod(indices, aux.size)]
+    ind2store = np.nonzero(np.any(weights, axis=0))
+    weights = weights[:, ind2store]
+    indices = indices[:, ind2store]
+    return weights, indices
+
+
+def imresizemex(inimg, weights, indices, dim):
+    in_shape = inimg.shape
+    w_shape = weights.shape
+    out_shape = list(in_shape)
+    out_shape[dim] = w_shape[0]
+    outimg = np.zeros(out_shape)
+    if dim == 0:
+        for i_img in range(in_shape[1]):
+            for i_w in range(w_shape[0]):
+                w = weights[i_w, :]
+                ind = indices[i_w, :]
+                im_slice = inimg[ind, i_img].astype(np.float64)
+                outimg[i_w, i_img] = np.sum(np.multiply(np.squeeze(im_slice, axis=0), w.T), axis=0)
+    elif dim == 1:
+        for i_img in range(in_shape[0]):
+            for i_w in range(w_shape[0]):
+                w = weights[i_w, :]
+                ind = indices[i_w, :]
+                im_slice = inimg[i_img, ind].astype(np.float64)
+                outimg[i_img, i_w] = np.sum(np.multiply(np.squeeze(im_slice, axis=0), w.T), axis=0)
+    if inimg.dtype == np.uint8:
+        outimg = np.clip(outimg, 0, 255)
+        return np.around(outimg).astype(np.uint8)
+    else:
+        return outimg
+
+
+def imresizevec(inimg, weights, indices, dim):
+    wshape = weights.shape
+    if dim == 0:
+        weights = weights.reshape((wshape[0], wshape[2], 1, 1))
+        outimg = np.sum(weights * ((inimg[indices].squeeze(axis=1)).astype(np.float64)), axis=1)
+    elif dim == 1:
+        weights = weights.reshape((1, wshape[0], wshape[2], 1))
+        outimg = np.sum(weights * ((inimg[:, indices].squeeze(axis=2)).astype(np.float64)), axis=2)
+    if inimg.dtype == np.uint8:
+        outimg = np.clip(outimg, 0, 255)
+        return np.around(outimg).astype(np.uint8)
+    else:
+        return outimg
+
+
+def resizeAlongDim(A, dim, weights, indices, mode="vec"):
+    if mode == "org":
+        out = imresizemex(A, weights, indices, dim)
+    else:
+        out = imresizevec(A, weights, indices, dim)
+    return out
+
+
+def imresize(I, scalar_scale=None, method='bicubic', output_shape=None, mode="vec"):
+    if method is 'bicubic':
+        kernel = cubic
+    elif method is 'bilinear':
+        kernel = triangle
+    else:
+        print('Error: Unidentified method supplied')
+
+    kernel_width = 4.0
+    # Fill scale and output_size
+    if scalar_scale is not None:
+        scalar_scale = float(scalar_scale)
+        scale = [scalar_scale, scalar_scale]
+        output_size = deriveSizeFromScale(I.shape, scale)
+    elif output_shape is not None:
+        scale = deriveScaleFromSize(I.shape, output_shape)
+        output_size = list(output_shape)
+    else:
+        print('Error: scalar_scale OR output_shape should be defined!')
+        return
+    scale_np = np.array(scale)
+    order = np.argsort(scale_np)
+    weights = []
+    indices = []
+    for k in range(2):
+        w, ind = contributions(I.shape[k], output_size[k], scale[k], kernel, kernel_width)
+        weights.append(w)
+        indices.append(ind)
+    B = np.copy(I)
+    flag2D = False
+    if B.ndim == 2:
+        B = np.expand_dims(B, axis=2)
+        flag2D = True
+    for k in range(2):
+        dim = order[k]
+        B = resizeAlongDim(B, dim, weights[dim], indices[dim], mode)
+    if flag2D:
+        B = np.squeeze(B, axis=2)
+    return B
+
+
+def convertDouble2Byte(I):
+    B = np.clip(I, 0.0, 1.0)
+    B = 255 * B
+    return np.around(B).astype(np.uint8)

models/SRFlow/code/models/SRFlow_model.py

@@ -0,0 +1,280 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import logging
+import sys
+sys.path.append('../..')
+from collections import OrderedDict
+from SRFlow.code.utils.util import get_resume_paths, opt_get
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+import SRFlow.code.models.networks as networks
+import SRFlow.code.models.lr_scheduler as lr_scheduler
+from .base_model import BaseModel
+
+logger = logging.getLogger('base')
+
+
+class SRFlowModel(BaseModel):
+    def __init__(self, opt, step):
+        super(SRFlowModel, self).__init__(opt)
+        self.opt = opt
+
+        self.heats = opt['val']['heats']
+        self.n_sample = opt['val']['n_sample']
+        self.hr_size = opt_get(opt, ['datasets', 'train', 'center_crop_hr_size'])
+        self.hr_size = 256 if self.hr_size is None else self.hr_size
+        self.lr_size = self.hr_size // opt['scale']
+
+        if opt['dist']:
+            self.rank = torch.distributed.get_rank()
+        else:
+            self.rank = -1  # non dist training
+        train_opt = opt['train']
+
+        # define network and load pretrained models
+        self.netG = networks.define_Flow(opt, step).to(self.device)
+        if opt['dist']:
+            self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
+        else:
+            self.netG = DataParallel(self.netG)
+        # print network
+        self.print_network()
+
+        if opt_get(opt, ['path', 'resume_state'], 1) is not None:
+            self.load()
+        else:
+            print("WARNING: skipping initial loading, due to resume_state None")
+
+        if self.is_train:
+            self.netG.train()
+
+            self.init_optimizer_and_scheduler(train_opt)
+            self.log_dict = OrderedDict()
+
+    def to(self, device):
+        self.device = device
+        self.netG.to(device)
+
+    def init_optimizer_and_scheduler(self, train_opt):
+        # optimizers
+        self.optimizers = []
+        wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+        optim_params_RRDB = []
+        optim_params_other = []
+        for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+            print(k, v.requires_grad)
+            if v.requires_grad:
+                if '.RRDB.' in k:
+                    optim_params_RRDB.append(v)
+                    print('opt', k)
+                else:
+                    optim_params_other.append(v)
+                if self.rank <= 0:
+                    logger.warning('Params [{:s}] will not optimize.'.format(k))
+
+        print('rrdb params', len(optim_params_RRDB))
+
+        self.optimizer_G = torch.optim.Adam(
+            [
+                {"params": optim_params_other, "lr": train_opt['lr_G'], 'beta1': train_opt['beta1'],
+                 'beta2': train_opt['beta2'], 'weight_decay': wd_G},
+                {"params": optim_params_RRDB, "lr": train_opt.get('lr_RRDB', train_opt['lr_G']),
+                 'beta1': train_opt['beta1'],
+                 'beta2': train_opt['beta2'], 'weight_decay': wd_G}
+            ],
+        )
+
+        self.optimizers.append(self.optimizer_G)
+        # schedulers
+        if train_opt['lr_scheme'] == 'MultiStepLR':
+            for optimizer in self.optimizers:
+                self.schedulers.append(
+                    lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
+                                                     restarts=train_opt['restarts'],
+                                                     weights=train_opt['restart_weights'],
+                                                     gamma=train_opt['lr_gamma'],
+                                                     clear_state=train_opt['clear_state'],
+                                                     lr_steps_invese=train_opt.get('lr_steps_inverse', [])))
+        elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
+            for optimizer in self.optimizers:
+                self.schedulers.append(
+                    lr_scheduler.CosineAnnealingLR_Restart(
+                        optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
+                        restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
+        else:
+            raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+    def add_optimizer_and_scheduler_RRDB(self, train_opt):
+        # optimizers
+        assert len(self.optimizers) == 1, self.optimizers
+        assert len(self.optimizer_G.param_groups[1]['params']) == 0, self.optimizer_G.param_groups[1]
+        for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+            if v.requires_grad:
+                if '.RRDB.' in k:
+                    self.optimizer_G.param_groups[1]['params'].append(v)
+        assert len(self.optimizer_G.param_groups[1]['params']) > 0
+
+    def feed_data(self, data, need_GT=True):
+        self.var_L = data['LQ'].to(self.device)  # LQ
+        if need_GT:
+            self.real_H = data['GT'].to(self.device)  # GT
+
+    def optimize_parameters(self, step):
+
+        train_RRDB_delay = opt_get(self.opt, ['network_G', 'train_RRDB_delay'])
+        if train_RRDB_delay is not None and step > int(train_RRDB_delay * self.opt['train']['niter']) \
+                and not self.netG.module.RRDB_training:
+            if self.netG.module.set_rrdb_training(True):
+                self.add_optimizer_and_scheduler_RRDB(self.opt['train'])
+
+        # self.print_rrdb_state()
+
+        self.netG.train()
+        self.log_dict = OrderedDict()
+        self.optimizer_G.zero_grad()
+
+        losses = {}
+        weight_fl = opt_get(self.opt, ['train', 'weight_fl'])
+        weight_fl = 1 if weight_fl is None else weight_fl
+        if weight_fl > 0:
+            z, nll, y_logits = self.netG(gt=self.real_H, lr=self.var_L, reverse=False)
+            nll_loss = torch.mean(nll)
+            losses['nll_loss'] = nll_loss * weight_fl
+
+        weight_l1 = opt_get(self.opt, ['train', 'weight_l1']) or 0
+        if weight_l1 > 0:
+            z = self.get_z(heat=0, seed=None, batch_size=self.var_L.shape[0], lr_shape=self.var_L.shape)
+            sr, logdet = self.netG(lr=self.var_L, z=z, eps_std=0, reverse=True, reverse_with_grad=True)
+            l1_loss = (sr - self.real_H).abs().mean()
+            losses['l1_loss'] = l1_loss * weight_l1
+
+        total_loss = sum(losses.values())
+        total_loss.backward()
+        self.optimizer_G.step()
+
+        mean = total_loss.item()
+        return mean
+
+    def print_rrdb_state(self):
+        for name, param in self.netG.module.named_parameters():
+            if "RRDB.conv_first.weight" in name:
+                print(name, param.requires_grad, param.data.abs().sum())
+        print('params', [len(p['params']) for p in self.optimizer_G.param_groups])
+
+    def test(self):
+        self.netG.eval()
+        self.fake_H = {}
+        for heat in self.heats:
+            for i in range(self.n_sample):
+                z = self.get_z(heat, seed=None, batch_size=self.var_L.shape[0], lr_shape=self.var_L.shape)
+                with torch.no_grad():
+                    self.fake_H[(heat, i)], logdet = self.netG(lr=self.var_L, z=z, eps_std=heat, reverse=True)
+        with torch.no_grad():
+            _, nll, _ = self.netG(gt=self.real_H, lr=self.var_L, reverse=False)
+        self.netG.train()
+        return nll.mean().item()
+
+    def get_encode_nll(self, lq, gt):
+        self.netG.eval()
+        with torch.no_grad():
+            _, nll, _ = self.netG(gt=gt, lr=lq, reverse=False)
+        self.netG.train()
+        return nll.mean().item()
+
+    def get_sr(self, lq, heat=None, seed=None, z=None, epses=None):
+        return self.get_sr_with_z(lq, heat, seed, z, epses)[0]
+
+    def get_encode_z(self, lq, gt, epses=None, add_gt_noise=True):
+        self.netG.eval()
+        with torch.no_grad():
+            z, _, _ = self.netG(gt=gt, lr=lq, reverse=False, epses=epses, add_gt_noise=add_gt_noise)
+        self.netG.train()
+        return z
+
+    def get_encode_z_and_nll(self, lq, gt, epses=None, add_gt_noise=True):
+        self.netG.eval()
+        with torch.no_grad():
+            z, nll, _ = self.netG(gt=gt, lr=lq, reverse=False, epses=epses, add_gt_noise=add_gt_noise)
+        self.netG.train()
+        return z, nll
+
+    def get_sr_with_z(self, lq, heat=None, seed=None, z=None, epses=None):
+        self.netG.eval()
+
+        z = self.get_z(heat, seed, batch_size=lq.shape[0], lr_shape=lq.shape) if z is None and epses is None else z
+
+        with torch.no_grad():
+            sr, logdet = self.netG(lr=lq, z=z, eps_std=heat, reverse=True, epses=epses)
+        self.netG.train()
+        return sr, z
+
+    def get_z(self, heat, seed=None, batch_size=1, lr_shape=None):
+        if seed: torch.manual_seed(seed)
+        if opt_get(self.opt, ['network_G', 'flow', 'split', 'enable']):
+            C = self.netG.module.flowUpsamplerNet.C
+            H = int(self.opt['scale'] * lr_shape[2] // self.netG.module.flowUpsamplerNet.scaleH)
+            W = int(self.opt['scale'] * lr_shape[3] // self.netG.module.flowUpsamplerNet.scaleW)
+            z = torch.normal(mean=0, std=heat, size=(batch_size, C, H, W)) if heat > 0 else torch.zeros(
+                (batch_size, C, H, W))
+        else:
+            L = opt_get(self.opt, ['network_G', 'flow', 'L']) or 3
+            fac = 2 ** (L - 3)
+            z_size = int(self.lr_size // (2 ** (L - 3)))
+            z = torch.normal(mean=0, std=heat, size=(batch_size, 3 * 8 * 8 * fac * fac, z_size, z_size))
+        return z
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self, need_GT=True):
+        out_dict = OrderedDict()
+        out_dict['LQ'] = self.var_L.detach()[0].float().cpu()
+        for heat in self.heats:
+            for i in range(self.n_sample):
+                out_dict[('SR', heat, i)] = self.fake_H[(heat, i)].detach()[0].float().cpu()
+        if need_GT:
+            out_dict['GT'] = self.real_H.detach()[0].float().cpu()
+        return out_dict
+
+    def print_network(self):
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel) or isinstance(self.netG, DistributedDataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+        if self.rank <= 0:
+            logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+            logger.info(s)
+
+    def load(self):
+        _, get_resume_model_path = get_resume_paths(self.opt)
+        if get_resume_model_path is not None:
+            self.load_network(get_resume_model_path, self.netG, strict=True, submodule=None)
+            return
+
+        load_path_G = self.opt['path']['pretrain_model_G']
+        load_submodule = self.opt['path']['load_submodule'] if 'load_submodule' in self.opt['path'].keys() else 'RRDB'
+        if load_path_G is not None:
+            logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG, self.opt['path'].get('strict_load', True),
+                              submodule=load_submodule)
+
+    def save(self, iter_label):
+        self.save_network(self.netG, 'G', iter_label)

models/SRFlow/code/models/SR_model.py

@@ -0,0 +1,217 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+import models.networks as networks
+import models.lr_scheduler as lr_scheduler
+from utils.util import opt_get
+from .base_model import BaseModel
+from models.modules.loss import CharbonnierLoss
+
+logger = logging.getLogger('base')
+
+
+class SRModel(BaseModel):
+    def __init__(self, opt, step):
+        super(SRModel, self).__init__(opt)
+
+        self.step = step
+
+        if opt['dist']:
+            self.rank = torch.distributed.get_rank()
+        else:
+            self.rank = -1  # non dist training
+        train_opt = opt['train']
+
+        # define network and load pretrained_models models
+        self.netG = networks.define_G(opt).to(self.device)
+        if opt['dist']:
+            self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
+        else:
+            self.netG = DataParallel(self.netG)
+        # print network
+        self.print_network()
+        self.load()
+
+        if self.is_train:
+            self.netG.train()
+
+            # loss
+            loss_type = train_opt['pixel_criterion']
+            if loss_type == 'l1':
+                self.cri_pix = nn.L1Loss().to(self.device)
+            elif loss_type == 'l2':
+                self.cri_pix = nn.MSELoss().to(self.device)
+            elif loss_type == 'cb':
+                self.cri_pix = CharbonnierLoss().to(self.device)
+            else:
+                raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type))
+            self.l_pix_w = train_opt['pixel_weight']
+
+            # optimizers
+            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+            optim_params = []
+            for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+                if v.requires_grad:
+                    optim_params.append(v)
+                else:
+                    if self.rank <= 0:
+                        logger.warning('Params [{:s}] will not optimize.'.format(k))
+            self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'],
+                                                weight_decay=wd_G,
+                                                betas=(train_opt['beta1'], train_opt['beta2']))
+            self.optimizers.append(self.optimizer_G)
+
+            # schedulers
+            if train_opt['lr_scheme'] == 'MultiStepLR':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(
+                        lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
+                                                         restarts=train_opt['restarts'],
+                                                         weights=train_opt['restart_weights'],
+                                                         gamma=train_opt['lr_gamma'],
+                                                         clear_state=train_opt['clear_state']))
+            elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(
+                        lr_scheduler.CosineAnnealingLR_Restart(
+                            optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
+                            restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
+            else:
+                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+            self.log_dict = OrderedDict()
+
+    def feed_data(self, data, need_GT=True):
+        self.var_L = data['LQ'].to(self.device)  # LQ
+        if need_GT:
+            self.real_H = data['GT'].to(self.device)  # GT
+
+    def to(self, device):
+        self.device = device
+        self.netG.to(device)
+
+    def optimize_parameters(self, step):
+        def getEnv(name): import os; return True if name in os.environ.keys() else False
+
+        if getEnv("DEBUG_FEED_IMAGES"):
+            import imageio
+            import random
+            i = random.randint(0, 10000)
+            label = self.var_L.cpu().numpy()[0].transpose([1, 2, 0])
+            print("var_L", label.min(), label.max(), label.shape)
+            imageio.imwrite("/tmp/{}_l.png".format(i), label)
+            image = self.real_H.cpu().numpy()[0].transpose([1, 2, 0])
+            print("self.real_H", image.min(), image.max(), image.shape)
+            imageio.imwrite("/tmp/{}_gt.png".format(i), image)
+        self.optimizer_G.zero_grad()
+        self.fake_H = self.netG(self.var_L)
+        l_pix = self.l_pix_w * self.cri_pix(self.fake_H, self.real_H.to(self.fake_H.device))
+        l_pix.backward()
+        self.optimizer_G.step()
+
+        # set log
+        self.log_dict['l_pix'] = l_pix.item()
+
+    def test(self):
+        self.netG.eval()
+        with torch.no_grad():
+            self.fake_H = self.netG(self.var_L)
+        self.netG.train()
+
+    def get_encode_nll(self, lq, gt):
+        return torch.ones(1) * 1e14
+
+    def get_sr(self, lq, heat=None, seed=None):
+        self.netG.eval()
+        sr = self.netG(lq)
+        self.netG.train()
+        return sr
+
+    def test_x8(self):
+        # from https://github.com/thstkdgus35/EDSR-PyTorch
+        self.netG.eval()
+
+        def _transform(v, op):
+            # if self.precision != 'single': v = v.float()
+            v2np = v.data.cpu().numpy()
+            if op == 'v':
+                tfnp = v2np[:, :, :, ::-1].copy()
+            elif op == 'h':
+                tfnp = v2np[:, :, ::-1, :].copy()
+            elif op == 't':
+                tfnp = v2np.transpose((0, 1, 3, 2)).copy()
+
+            ret = torch.Tensor(tfnp).to(self.device)
+            # if self.precision == 'half': ret = ret.half()
+
+            return ret
+
+        lr_list = [self.var_L]
+        for tf in 'v', 'h', 't':
+            lr_list.extend([_transform(t, tf) for t in lr_list])
+        with torch.no_grad():
+            sr_list = [self.netG(aug) for aug in lr_list]
+        for i in range(len(sr_list)):
+            if i > 3:
+                sr_list[i] = _transform(sr_list[i], 't')
+            if i % 4 > 1:
+                sr_list[i] = _transform(sr_list[i], 'h')
+            if (i % 4) % 2 == 1:
+                sr_list[i] = _transform(sr_list[i], 'v')
+
+        output_cat = torch.cat(sr_list, dim=0)
+        self.fake_H = output_cat.mean(dim=0, keepdim=True)
+        self.netG.train()
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self, need_GT=True):
+        out_dict = OrderedDict()
+        out_dict['LQ'] = self.var_L.detach()[0].float().cpu()
+        out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
+        if need_GT:
+            out_dict['GT'] = self.real_H.detach()[0].float().cpu()
+        return out_dict
+
+    def print_network(self):
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel) or isinstance(self.netG, DistributedDataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+        if self.rank <= 0:
+            logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+            logger.info(s)
+
+    def load(self):
+        load_path_G = self.opt['path']['pretrain_model_G']
+        if load_path_G is not None:
+            logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG, self.opt['path']['strict_load'])
+
+    def save(self, iter_label):
+        self.save_network(self.netG, 'G', iter_label)
+
+    def get_encode_z_and_nll(self, *args, **kwargs):
+        return [], torch.zeros(1)

models/SRFlow/code/models/__init__.py

@@ -0,0 +1,52 @@
+import importlib
+import logging
+import os
+
+try:
+    import local_config
+except:
+    local_config = None
+
+
+logger = logging.getLogger('base')
+
+
+def find_model_using_name(model_name):
+    # Given the option --model [modelname],
+    # the file "models/modelname_model.py"
+    # will be imported.
+    model_filename = "SRFlow.code.models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+
+    # In the file, the class called ModelNameModel() will
+    # be instantiated. It has to be a subclass of torch.nn.Module,
+    # and it is case-insensitive.
+    model = None
+    target_model_name = model_name.replace('_', '') + 'Model'
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower():
+            model = cls
+
+    if model is None:
+        print(
+            "In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s." % (
+            model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+def create_model(opt, step=0, **opt_kwargs):
+    if local_config is not None:
+        opt['path']['pretrain_model_G'] = os.path.join(local_config.checkpoint_path, os.path.basename(opt['path']['results_root'] + '.pth'))
+
+    for k, v in opt_kwargs.items():
+        opt[k] = v
+
+    model = opt['model']
+
+    M = find_model_using_name(model)
+
+    m = M(opt, step)
+    logger.info('Model [{:s}] is created.'.format(m.__class__.__name__))
+    return m

models/SRFlow/code/models/base_model.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import os
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+import natsort
+import glob
+
+
+class BaseModel():
+    def __init__(self, opt):
+        self.opt = opt
+        self.device = torch.device('cuda' if opt.get('gpu_ids', None) is not None else 'cpu')
+        self.is_train = opt['is_train']
+        self.schedulers = []
+        self.optimizers = []
+
+    def feed_data(self, data):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        pass
+
+    def get_current_losses(self):
+        pass
+
+    def print_network(self):
+        pass
+
+    def save(self, label):
+        pass
+
+    def load(self):
+        pass
+
+    def _set_lr(self, lr_groups_l):
+        ''' set learning rate for warmup,
+        lr_groups_l: list for lr_groups. each for a optimizer'''
+        for optimizer, lr_groups in zip(self.optimizers, lr_groups_l):
+            for param_group, lr in zip(optimizer.param_groups, lr_groups):
+                param_group['lr'] = lr
+
+    def _get_init_lr(self):
+        # get the initial lr, which is set by the scheduler
+        init_lr_groups_l = []
+        for optimizer in self.optimizers:
+            init_lr_groups_l.append([v['initial_lr'] for v in optimizer.param_groups])
+        return init_lr_groups_l
+
+    def update_learning_rate(self, cur_iter, warmup_iter=-1):
+        for scheduler in self.schedulers:
+            scheduler.step()
+        #### set up warm up learning rate
+        if cur_iter < warmup_iter:
+            # get initial lr for each group
+            init_lr_g_l = self._get_init_lr()
+            # modify warming-up learning rates
+            warm_up_lr_l = []
+            for init_lr_g in init_lr_g_l:
+                warm_up_lr_l.append([v / warmup_iter * cur_iter for v in init_lr_g])
+            # set learning rate
+            self._set_lr(warm_up_lr_l)
+
+    def get_current_learning_rate(self):
+        # return self.schedulers[0].get_lr()[0]
+        return self.optimizers[0].param_groups[0]['lr']
+
+    def get_network_description(self, network):
+        '''Get the string and total parameters of the network'''
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        s = str(network)
+        n = sum(map(lambda x: x.numel(), network.parameters()))
+        return s, n
+
+    def save_network(self, network, network_label, iter_label):
+        paths = natsort.natsorted(glob.glob(os.path.join(self.opt['path']['models'], "*_{}.pth".format(network_label))),
+                                  reverse=True)
+        paths = [p for p in paths if
+                 "latest_" not in p and not any([str(i * 10000) in p.split("/")[-1].split("_") for i in range(101)])]
+        if len(paths) > 2:
+            for path in paths[2:]:
+                os.remove(path)
+        save_filename = '{}_{}.pth'.format(iter_label, network_label)
+        save_path = os.path.join(self.opt['path']['models'], save_filename)
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        state_dict = network.state_dict()
+        for key, param in state_dict.items():
+            state_dict[key] = param.cpu()
+        torch.save(state_dict, save_path)
+
+    def load_network(self, load_path, network, strict=True, submodule=None):
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        if not (submodule is None or submodule.lower() == 'none'.lower()):
+            network = network.__getattr__(submodule)
+        load_net = torch.load(load_path)
+        load_net_clean = OrderedDict()  # remove unnecessary 'module.'
+        for k, v in load_net.items():
+            if k.startswith('module.'):
+                load_net_clean[k[7:]] = v
+            else:
+                load_net_clean[k] = v
+        network.load_state_dict(load_net_clean, strict=strict)
+
+    def save_training_state(self, epoch, iter_step):
+        '''Saves training state during training, which will be used for resuming'''
+        state = {'epoch': epoch, 'iter': iter_step, 'schedulers': [], 'optimizers': []}
+        for s in self.schedulers:
+            state['schedulers'].append(s.state_dict())
+        for o in self.optimizers:
+            state['optimizers'].append(o.state_dict())
+        save_filename = '{}.state'.format(iter_step)
+        save_path = os.path.join(self.opt['path']['training_state'], save_filename)
+
+        paths = natsort.natsorted(glob.glob(os.path.join(self.opt['path']['training_state'], "*.state")),
+                                  reverse=True)
+        paths = [p for p in paths if "latest_" not in p]
+        if len(paths) > 2:
+            for path in paths[2:]:
+                os.remove(path)
+
+        torch.save(state, save_path)
+
+    def resume_training(self, resume_state):
+        '''Resume the optimizers and schedulers for training'''
+        resume_optimizers = resume_state['optimizers']
+        resume_schedulers = resume_state['schedulers']
+        assert len(resume_optimizers) == len(self.optimizers), 'Wrong lengths of optimizers'
+        assert len(resume_schedulers) == len(self.schedulers), 'Wrong lengths of schedulers'
+        for i, o in enumerate(resume_optimizers):
+            self.optimizers[i].load_state_dict(o)
+        for i, s in enumerate(resume_schedulers):
+            self.schedulers[i].load_state_dict(s)

models/SRFlow/code/models/lr_scheduler.py

@@ -0,0 +1,163 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import math
+from collections import Counter
+from collections import defaultdict
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class MultiStepLR_Restart(_LRScheduler):
+    def __init__(self, optimizer, milestones, restarts=None, weights=None, gamma=0.1,
+                 clear_state=False, last_epoch=-1, lr_steps_invese=None):
+        assert lr_steps_invese is not None, "Use empty list"
+        self.milestones = Counter(milestones)
+        self.lr_steps_inverse = Counter(lr_steps_invese)
+        self.gamma = gamma
+        self.clear_state = clear_state
+        self.restarts = restarts if restarts else [0]
+        self.restart_weights = weights if weights else [1]
+        assert len(self.restarts) == len(
+            self.restart_weights), 'restarts and their weights do not match.'
+        super(MultiStepLR_Restart, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch in self.restarts:
+            if self.clear_state:
+                self.optimizer.state = defaultdict(dict)
+            weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+            return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+        if self.last_epoch not in self.milestones and self.last_epoch not in self.lr_steps_inverse:
+            return [group['lr'] for group in self.optimizer.param_groups]
+        return [
+            group['lr'] * (self.gamma ** self.milestones[self.last_epoch]) *
+            (self.gamma ** (-self.lr_steps_inverse[self.last_epoch]))
+            for group in self.optimizer.param_groups
+        ]
+
+
+class CosineAnnealingLR_Restart(_LRScheduler):
+    def __init__(self, optimizer, T_period, restarts=None, weights=None, eta_min=0, last_epoch=-1):
+        self.T_period = T_period
+        self.T_max = self.T_period[0]  # current T period
+        self.eta_min = eta_min
+        self.restarts = restarts if restarts else [0]
+        self.restart_weights = weights if weights else [1]
+        self.last_restart = 0
+        assert len(self.restarts) == len(
+            self.restart_weights), 'restarts and their weights do not match.'
+        super(CosineAnnealingLR_Restart, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch == 0:
+            return self.base_lrs
+        elif self.last_epoch in self.restarts:
+            self.last_restart = self.last_epoch
+            self.T_max = self.T_period[self.restarts.index(self.last_epoch) + 1]
+            weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+            return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+        elif (self.last_epoch - self.last_restart - 1 - self.T_max) % (2 * self.T_max) == 0:
+            return [
+                group['lr'] + (base_lr - self.eta_min) * (1 - math.cos(math.pi / self.T_max)) / 2
+                for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
+            ]
+        return [(1 + math.cos(math.pi * (self.last_epoch - self.last_restart) / self.T_max)) /
+                (1 + math.cos(math.pi * ((self.last_epoch - self.last_restart) - 1) / self.T_max)) *
+                (group['lr'] - self.eta_min) + self.eta_min
+                for group in self.optimizer.param_groups]
+
+
+if __name__ == "__main__":
+    optimizer = torch.optim.Adam([torch.zeros(3, 64, 3, 3)], lr=2e-4, weight_decay=0,
+                                 betas=(0.9, 0.99))
+    ##############################
+    # MultiStepLR_Restart
+    ##############################
+    ## Original
+    lr_steps = [200000, 400000, 600000, 800000]
+    restarts = None
+    restart_weights = None
+
+    ## two
+    lr_steps = [100000, 200000, 300000, 400000, 490000, 600000, 700000, 800000, 900000, 990000]
+    restarts = [500000]
+    restart_weights = [1]
+
+    ## four
+    lr_steps = [
+        50000, 100000, 150000, 200000, 240000, 300000, 350000, 400000, 450000, 490000, 550000,
+        600000, 650000, 700000, 740000, 800000, 850000, 900000, 950000, 990000
+    ]
+    restarts = [250000, 500000, 750000]
+    restart_weights = [1, 1, 1]
+
+    scheduler = MultiStepLR_Restart(optimizer, lr_steps, restarts, restart_weights, gamma=0.5,
+                                    clear_state=False)
+
+    ##############################
+    # Cosine Annealing Restart
+    ##############################
+    ## two
+    T_period = [500000, 500000]
+    restarts = [500000]
+    restart_weights = [1]
+
+    ## four
+    T_period = [250000, 250000, 250000, 250000]
+    restarts = [250000, 500000, 750000]
+    restart_weights = [1, 1, 1]
+
+    scheduler = CosineAnnealingLR_Restart(optimizer, T_period, eta_min=1e-7, restarts=restarts,
+                                          weights=restart_weights)
+
+    ##############################
+    # Draw figure
+    ##############################
+    N_iter = 1000000
+    lr_l = list(range(N_iter))
+    for i in range(N_iter):
+        scheduler.step()
+        current_lr = optimizer.param_groups[0]['lr']
+        lr_l[i] = current_lr
+
+    import matplotlib as mpl
+    from matplotlib import pyplot as plt
+    import matplotlib.ticker as mtick
+
+    mpl.style.use('default')
+    import seaborn
+
+    seaborn.set(style='whitegrid')
+    seaborn.set_context('paper')
+
+    plt.figure(1)
+    plt.subplot(111)
+    plt.ticklabel_format(style='sci', axis='x', scilimits=(0, 0))
+    plt.title('Title', fontsize=16, color='k')
+    plt.plot(list(range(N_iter)), lr_l, linewidth=1.5, label='learning rate scheme')
+    legend = plt.legend(loc='upper right', shadow=False)
+    ax = plt.gca()
+    labels = ax.get_xticks().tolist()
+    for k, v in enumerate(labels):
+        labels[k] = str(int(v / 1000)) + 'K'
+    ax.set_xticklabels(labels)
+    ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))
+
+    ax.set_ylabel('Learning rate')
+    ax.set_xlabel('Iteration')
+    fig = plt.gcf()
+    plt.show()

models/SRFlow/code/models/modules/FlowActNorms.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+from torch import nn as nn
+import sys
+sys.path.append('../../../..')
+from SRFlow.code.models.modules import thops
+
+
+class _ActNorm(nn.Module):
+    """
+    Activation Normalization
+    Initialize the bias and scale with a given minibatch,
+    so that the output per-channel have zero mean and unit variance for that.
+
+    After initialization, `bias` and `logs` will be trained as parameters.
+    """
+
+    def __init__(self, num_features, scale=1.):
+        super().__init__()
+        # register mean and scale
+        size = [1, num_features, 1, 1]
+        self.register_parameter("bias", nn.Parameter(torch.zeros(*size)))
+        self.register_parameter("logs", nn.Parameter(torch.zeros(*size)))
+        self.num_features = num_features
+        self.scale = float(scale)
+        self.inited = False
+
+    def _check_input_dim(self, input):
+        return NotImplemented
+
+    def initialize_parameters(self, input):
+        self._check_input_dim(input)
+        if not self.training:
+            return
+        if (self.bias != 0).any():
+            self.inited = True
+            return
+        assert input.device == self.bias.device, (input.device, self.bias.device)
+        with torch.no_grad():
+            bias = thops.mean(input.clone(), dim=[0, 2, 3], keepdim=True) * -1.0
+            vars = thops.mean((input.clone() + bias) ** 2, dim=[0, 2, 3], keepdim=True)
+            logs = torch.log(self.scale / (torch.sqrt(vars) + 1e-6))
+            self.bias.data.copy_(bias.data)
+            self.logs.data.copy_(logs.data)
+            self.inited = True
+
+    def _center(self, input, reverse=False, offset=None):
+        bias = self.bias
+
+        if offset is not None:
+            bias = bias + offset
+
+        if not reverse:
+            return input + bias
+        else:
+            return input - bias
+
+    def _scale(self, input, logdet=None, reverse=False, offset=None):
+        logs = self.logs
+
+        if offset is not None:
+            logs = logs + offset
+
+        if not reverse:
+            input = input * torch.exp(logs) # should have shape batchsize, n_channels, 1, 1
+            # input = input * torch.exp(logs+logs_offset)
+        else:
+            input = input * torch.exp(-logs)
+        if logdet is not None:
+            """
+            logs is log_std of `mean of channels`
+            so we need to multiply pixels
+            """
+            dlogdet = thops.sum(logs) * thops.pixels(input)
+            if reverse:
+                dlogdet *= -1
+            logdet = logdet + dlogdet
+        return input, logdet
+
+    def forward(self, input, logdet=None, reverse=False, offset_mask=None, logs_offset=None, bias_offset=None):
+        if not self.inited:
+            self.initialize_parameters(input)
+        self._check_input_dim(input)
+
+        if offset_mask is not None:
+            logs_offset *= offset_mask
+            bias_offset *= offset_mask
+        # no need to permute dims as old version
+        if not reverse:
+            # center and scale
+
+            # self.input = input
+            input = self._center(input, reverse, bias_offset)
+            input, logdet = self._scale(input, logdet, reverse, logs_offset)
+        else:
+            # scale and center
+            input, logdet = self._scale(input, logdet, reverse, logs_offset)
+            input = self._center(input, reverse, bias_offset)
+        return input, logdet
+
+
+class ActNorm2d(_ActNorm):
+    def __init__(self, num_features, scale=1.):
+        super().__init__(num_features, scale)
+
+    def _check_input_dim(self, input):
+        assert len(input.size()) == 4
+        assert input.size(1) == self.num_features, (
+            "[ActNorm]: input should be in shape as `BCHW`,"
+            " channels should be {} rather than {}".format(
+                self.num_features, input.size()))
+
+
+class MaskedActNorm2d(ActNorm2d):
+    def __init__(self, num_features, scale=1.):
+        super().__init__(num_features, scale)
+
+    def forward(self, input, mask, logdet=None, reverse=False):
+
+        assert mask.dtype == torch.bool
+        output, logdet_out = super().forward(input, logdet, reverse)
+
+        input[mask] = output[mask]
+        logdet[mask] = logdet_out[mask]
+
+        return input, logdet
+

models/SRFlow/code/models/modules/FlowAffineCouplingsAblation.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+from torch import nn as nn
+import sys
+sys.path.append('../../../..')
+from SRFlow.code.models.modules import thops
+from SRFlow.code.models.modules.flow import Conv2d, Conv2dZeros
+from SRFlow.code.utils.util import opt_get
+
+
+class CondAffineSeparatedAndCond(nn.Module):
+    def __init__(self, in_channels, opt):
+        super().__init__()
+        self.need_features = True
+        self.in_channels = in_channels
+        self.in_channels_rrdb = 320
+        self.kernel_hidden = 1
+        self.affine_eps = 0.0001
+        self.n_hidden_layers = 1
+        hidden_channels = opt_get(opt, ['network_G', 'flow', 'CondAffineSeparatedAndCond', 'hidden_channels'])
+        self.hidden_channels = 64 if hidden_channels is None else hidden_channels
+
+        self.affine_eps = opt_get(opt, ['network_G', 'flow', 'CondAffineSeparatedAndCond', 'eps'],  0.0001)
+
+        self.channels_for_nn = self.in_channels // 2
+        self.channels_for_co = self.in_channels - self.channels_for_nn
+
+        if self.channels_for_nn is None:
+            self.channels_for_nn = self.in_channels // 2
+
+        self.fAffine = self.F(in_channels=self.channels_for_nn + self.in_channels_rrdb,
+                              out_channels=self.channels_for_co * 2,
+                              hidden_channels=self.hidden_channels,
+                              kernel_hidden=self.kernel_hidden,
+                              n_hidden_layers=self.n_hidden_layers)
+
+        self.fFeatures = self.F(in_channels=self.in_channels_rrdb,
+                                out_channels=self.in_channels * 2,
+                                hidden_channels=self.hidden_channels,
+                                kernel_hidden=self.kernel_hidden,
+                                n_hidden_layers=self.n_hidden_layers)
+
+    def forward(self, input: torch.Tensor, logdet=None, reverse=False, ft=None):
+        if not reverse:
+            z = input
+            assert z.shape[1] == self.in_channels, (z.shape[1], self.in_channels)
+
+            # Feature Conditional
+            scaleFt, shiftFt = self.feature_extract(ft, self.fFeatures)
+            z = z + shiftFt
+            z = z * scaleFt
+            logdet = logdet + self.get_logdet(scaleFt)
+
+            # Self Conditional
+            z1, z2 = self.split(z)
+            scale, shift = self.feature_extract_aff(z1, ft, self.fAffine)
+            self.asserts(scale, shift, z1, z2)
+            z2 = z2 + shift
+            z2 = z2 * scale
+
+            logdet = logdet + self.get_logdet(scale)
+            z = thops.cat_feature(z1, z2)
+            output = z
+        else:
+            z = input
+
+            # Self Conditional
+            z1, z2 = self.split(z)
+            scale, shift = self.feature_extract_aff(z1, ft, self.fAffine)
+            self.asserts(scale, shift, z1, z2)
+            z2 = z2 / scale
+            z2 = z2 - shift
+            z = thops.cat_feature(z1, z2)
+            logdet = logdet - self.get_logdet(scale)
+
+            # Feature Conditional
+            scaleFt, shiftFt = self.feature_extract(ft, self.fFeatures)
+            z = z / scaleFt
+            z = z - shiftFt
+            logdet = logdet - self.get_logdet(scaleFt)
+
+            output = z
+        return output, logdet
+
+    def asserts(self, scale, shift, z1, z2):
+        assert z1.shape[1] == self.channels_for_nn, (z1.shape[1], self.channels_for_nn)
+        assert z2.shape[1] == self.channels_for_co, (z2.shape[1], self.channels_for_co)
+        assert scale.shape[1] == shift.shape[1], (scale.shape[1], shift.shape[1])
+        assert scale.shape[1] == z2.shape[1], (scale.shape[1], z1.shape[1], z2.shape[1])
+
+    def get_logdet(self, scale):
+        return thops.sum(torch.log(scale), dim=[1, 2, 3])
+
+    def feature_extract(self, z, f):
+        h = f(z)
+        shift, scale = thops.split_feature(h, "cross")
+        scale = (torch.sigmoid(scale + 2.) + self.affine_eps)
+        return scale, shift
+
+    def feature_extract_aff(self, z1, ft, f):
+        z = torch.cat([z1, ft], dim=1)
+        h = f(z)
+        shift, scale = thops.split_feature(h, "cross")
+        scale = (torch.sigmoid(scale + 2.) + self.affine_eps)
+        return scale, shift
+
+    def split(self, z):
+        z1 = z[:, :self.channels_for_nn]
+        z2 = z[:, self.channels_for_nn:]
+        assert z1.shape[1] + z2.shape[1] == z.shape[1], (z1.shape[1], z2.shape[1], z.shape[1])
+        return z1, z2
+
+    def F(self, in_channels, out_channels, hidden_channels, kernel_hidden=1, n_hidden_layers=1):
+        layers = [Conv2d(in_channels, hidden_channels), nn.ReLU(inplace=False)]
+
+        for _ in range(n_hidden_layers):
+            layers.append(Conv2d(hidden_channels, hidden_channels, kernel_size=[kernel_hidden, kernel_hidden]))
+            layers.append(nn.ReLU(inplace=False))
+        layers.append(Conv2dZeros(hidden_channels, out_channels))
+
+        return nn.Sequential(*layers)

models/SRFlow/code/models/modules/FlowStep.py

@@ -0,0 +1,139 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+from torch import nn as nn
+import sys
+sys.path.append('../../../..')
+import SRFlow.code.models.modules
+from SRFlow.code.models.modules.Permutations import InvertibleConv1x1
+from SRFlow.code.models.modules.FlowActNorms import ActNorm2d
+from SRFlow.code.models.modules import flow, thops, FlowAffineCouplingsAblation
+from SRFlow.code.utils.util import opt_get
+
+
+def getConditional(rrdbResults, position):
+    img_ft = rrdbResults if isinstance(rrdbResults, torch.Tensor) else rrdbResults[position]
+    return img_ft
+
+
+class FlowStep(nn.Module):
+    FlowPermutation = {
+        "reverse": lambda obj, z, logdet, rev: (obj.reverse(z, rev), logdet),
+        "shuffle": lambda obj, z, logdet, rev: (obj.shuffle(z, rev), logdet),
+        "invconv": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "squeeze_invconv": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "resqueeze_invconv_alternating_2_3": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "resqueeze_invconv_3": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "InvertibleConv1x1GridAlign": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "InvertibleConv1x1SubblocksShuf": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "InvertibleConv1x1GridAlignIndepBorder": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+        "InvertibleConv1x1GridAlignIndepBorder4": lambda obj, z, logdet, rev: obj.invconv(z, logdet, rev),
+    }
+
+    def __init__(self, in_channels, hidden_channels,
+                 actnorm_scale=1.0, flow_permutation="invconv", flow_coupling="additive",
+                 LU_decomposed=False, opt=None, image_injector=None, idx=None, acOpt=None, normOpt=None, in_shape=None,
+                 position=None):
+        # check configures
+        assert flow_permutation in FlowStep.FlowPermutation, \
+            "float_permutation should be in `{}`".format(
+                FlowStep.FlowPermutation.keys())
+        super().__init__()
+        self.flow_permutation = flow_permutation
+        self.flow_coupling = flow_coupling
+        self.image_injector = image_injector
+
+        self.norm_type = normOpt['type'] if normOpt else 'ActNorm2d'
+        self.position = normOpt['position'] if normOpt else None
+
+        self.in_shape = in_shape
+        self.position = position
+        self.acOpt = acOpt
+
+        # 1. actnorm
+        self.actnorm = ActNorm2d(in_channels, actnorm_scale)
+
+        # 2. permute
+        if flow_permutation == "invconv":
+            self.invconv = InvertibleConv1x1(
+                in_channels, LU_decomposed=LU_decomposed)
+
+        # 3. coupling
+        if flow_coupling == "CondAffineSeparatedAndCond":
+            self.affine = FlowAffineCouplingsAblation.CondAffineSeparatedAndCond(in_channels=in_channels,
+                                                                                                opt=opt)
+        elif flow_coupling == "noCoupling":
+            pass
+        else:
+            raise RuntimeError("coupling not Found:", flow_coupling)
+
+    def forward(self, input, logdet=None, reverse=False, rrdbResults=None):
+        if not reverse:
+            return self.normal_flow(input, logdet, rrdbResults)
+        else:
+            return self.reverse_flow(input, logdet, rrdbResults)
+
+    def normal_flow(self, z, logdet, rrdbResults=None):
+        if self.flow_coupling == "bentIdentityPreAct":
+            z, logdet = self.bentIdentPar(z, logdet, reverse=False)
+
+        # 1. actnorm
+        if self.norm_type == "ConditionalActNormImageInjector":
+            img_ft = getConditional(rrdbResults, self.position)
+            z, logdet = self.actnorm(z, img_ft=img_ft, logdet=logdet, reverse=False)
+        elif self.norm_type == "noNorm":
+            pass
+        else:
+            z, logdet = self.actnorm(z, logdet=logdet, reverse=False)
+
+        # 2. permute
+        z, logdet = FlowStep.FlowPermutation[self.flow_permutation](
+            self, z, logdet, False)
+
+        need_features = self.affine_need_features()
+
+        # 3. coupling
+        if need_features or self.flow_coupling in ["condAffine", "condFtAffine", "condNormAffine"]:
+            img_ft = getConditional(rrdbResults, self.position)
+            z, logdet = self.affine(input=z, logdet=logdet, reverse=False, ft=img_ft)
+        return z, logdet
+
+    def reverse_flow(self, z, logdet, rrdbResults=None):
+
+        need_features = self.affine_need_features()
+
+        # 1.coupling
+        if need_features or self.flow_coupling in ["condAffine", "condFtAffine", "condNormAffine"]:
+            img_ft = getConditional(rrdbResults, self.position)
+            z, logdet = self.affine(input=z, logdet=logdet, reverse=True, ft=img_ft)
+
+        # 2. permute
+        z, logdet = FlowStep.FlowPermutation[self.flow_permutation](
+            self, z, logdet, True)
+
+        # 3. actnorm
+        z, logdet = self.actnorm(z, logdet=logdet, reverse=True)
+
+        return z, logdet
+
+    def affine_need_features(self):
+        need_features = False
+        try:
+            need_features = self.affine.need_features
+        except:
+            pass
+        return need_features

models/SRFlow/code/models/modules/FlowUpsamplerNet.py

@@ -0,0 +1,310 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import numpy as np
+import torch
+from torch import nn as nn
+import sys
+sys.path.append('../../../..')
+import SRFlow.code.models.modules.Split
+from SRFlow.code.models.modules import flow, thops
+from SRFlow.code.models.modules.Split import Split2d
+from SRFlow.code.models.modules.glow_arch import f_conv2d_bias
+from SRFlow.code.models.modules.FlowStep import FlowStep
+from SRFlow.code.utils.util import opt_get
+
+
+class FlowUpsamplerNet(nn.Module):
+    def __init__(self, image_shape, hidden_channels, K, L=None,
+                 actnorm_scale=1.0,
+                 flow_permutation=None,
+                 flow_coupling="affine",
+                 LU_decomposed=False, opt=None):
+
+        super().__init__()
+
+        self.layers = nn.ModuleList()
+        self.output_shapes = []
+        self.L = opt_get(opt, ['network_G', 'flow', 'L'])
+        self.K = opt_get(opt, ['network_G', 'flow', 'K'])
+        if isinstance(self.K, int):
+            self.K = [K for K in [K, ] * (self.L + 1)]
+
+        self.opt = opt
+        H, W, self.C = image_shape
+        self.check_image_shape()
+
+        if opt['scale'] == 16:
+            self.levelToName = {
+                0: 'fea_up16',
+                1: 'fea_up8',
+                2: 'fea_up4',
+                3: 'fea_up2',
+                4: 'fea_up1',
+            }
+
+        if opt['scale'] == 8:
+            self.levelToName = {
+                0: 'fea_up8',
+                1: 'fea_up4',
+                2: 'fea_up2',
+                3: 'fea_up1',
+                4: 'fea_up0'
+            }
+
+        elif opt['scale'] == 4:
+            self.levelToName = {
+                0: 'fea_up4',
+                1: 'fea_up2',
+                2: 'fea_up1',
+                3: 'fea_up0',
+                4: 'fea_up-1'
+            }
+
+        affineInCh = self.get_affineInCh(opt_get)
+        flow_permutation = self.get_flow_permutation(flow_permutation, opt)
+
+        normOpt = opt_get(opt, ['network_G', 'flow', 'norm'])
+
+        conditional_channels = {}
+        n_rrdb = self.get_n_rrdb_channels(opt, opt_get)
+        n_bypass_channels = opt_get(opt, ['network_G', 'flow', 'levelConditional', 'n_channels'])
+        conditional_channels[0] = n_rrdb
+        for level in range(1, self.L + 1):
+            # Level 1 gets conditionals from 2, 3, 4 => L - level
+            # Level 2 gets conditionals from 3, 4
+            # Level 3 gets conditionals from 4
+            # Level 4 gets conditionals from None
+            n_bypass = 0 if n_bypass_channels is None else (self.L - level) * n_bypass_channels
+            conditional_channels[level] = n_rrdb + n_bypass
+
+        # Upsampler
+        for level in range(1, self.L + 1):
+            # 1. Squeeze
+            H, W = self.arch_squeeze(H, W)
+
+            # 2. K FlowStep
+            self.arch_additionalFlowAffine(H, LU_decomposed, W, actnorm_scale, hidden_channels, opt)
+            self.arch_FlowStep(H, self.K[level], LU_decomposed, W, actnorm_scale, affineInCh, flow_coupling,
+                               flow_permutation,
+                               hidden_channels, normOpt, opt, opt_get,
+                               n_conditinal_channels=conditional_channels[level])
+            # Split
+            self.arch_split(H, W, level, self.L, opt, opt_get)
+
+        if opt_get(opt, ['network_G', 'flow', 'split', 'enable']):
+            self.f = f_conv2d_bias(affineInCh, 2 * 3 * 64 // 2 // 2)
+        else:
+            self.f = f_conv2d_bias(affineInCh, 2 * 3 * 64)
+
+        self.H = H
+        self.W = W
+        self.scaleH = 160 / H
+        self.scaleW = 160 / W
+
+    def get_n_rrdb_channels(self, opt, opt_get):
+        blocks = opt_get(opt, ['network_G', 'flow', 'stackRRDB', 'blocks'])
+        n_rrdb = 64 if blocks is None else (len(blocks) + 1) * 64
+        return n_rrdb
+
+    def arch_FlowStep(self, H, K, LU_decomposed, W, actnorm_scale, affineInCh, flow_coupling, flow_permutation,
+                      hidden_channels, normOpt, opt, opt_get, n_conditinal_channels=None):
+        condAff = self.get_condAffSetting(opt, opt_get)
+        if condAff is not None:
+            condAff['in_channels_rrdb'] = n_conditinal_channels
+
+        for k in range(K):
+            position_name = get_position_name(H, self.opt['scale'])
+            if normOpt: normOpt['position'] = position_name
+
+            self.layers.append(
+                FlowStep(in_channels=self.C,
+                         hidden_channels=hidden_channels,
+                         actnorm_scale=actnorm_scale,
+                         flow_permutation=flow_permutation,
+                         flow_coupling=flow_coupling,
+                         acOpt=condAff,
+                         position=position_name,
+                         LU_decomposed=LU_decomposed, opt=opt, idx=k, normOpt=normOpt))
+            self.output_shapes.append(
+                [-1, self.C, H, W])
+
+    def get_condAffSetting(self, opt, opt_get):
+        condAff = opt_get(opt, ['network_G', 'flow', 'condAff']) or None
+        condAff = opt_get(opt, ['network_G', 'flow', 'condFtAffine']) or condAff
+        return condAff
+
+    def arch_split(self, H, W, L, levels, opt, opt_get):
+        correct_splits = opt_get(opt, ['network_G', 'flow', 'split', 'correct_splits'], False)
+        correction = 0 if correct_splits else 1
+        if opt_get(opt, ['network_G', 'flow', 'split', 'enable']) and L < levels - correction:
+            logs_eps = opt_get(opt, ['network_G', 'flow', 'split', 'logs_eps']) or 0
+            consume_ratio = opt_get(opt, ['network_G', 'flow', 'split', 'consume_ratio']) or 0.5
+            position_name = get_position_name(H, self.opt['scale'])
+            position = position_name if opt_get(opt, ['network_G', 'flow', 'split', 'conditional']) else None
+            cond_channels = opt_get(opt, ['network_G', 'flow', 'split', 'cond_channels'])
+            cond_channels = 0 if cond_channels is None else cond_channels
+
+            t = opt_get(opt, ['network_G', 'flow', 'split', 'type'], 'Split2d')
+
+            if t == 'Split2d':
+                split = Split2d(num_channels=self.C, logs_eps=logs_eps, position=position,
+                                                     cond_channels=cond_channels, consume_ratio=consume_ratio, opt=opt)
+            self.layers.append(split)
+            self.output_shapes.append([-1, split.num_channels_pass, H, W])
+            self.C = split.num_channels_pass
+
+    def arch_additionalFlowAffine(self, H, LU_decomposed, W, actnorm_scale, hidden_channels, opt):
+        if 'additionalFlowNoAffine' in opt['network_G']['flow']:
+            n_additionalFlowNoAffine = int(opt['network_G']['flow']['additionalFlowNoAffine'])
+            for _ in range(n_additionalFlowNoAffine):
+                self.layers.append(
+                    FlowStep(in_channels=self.C,
+                             hidden_channels=hidden_channels,
+                             actnorm_scale=actnorm_scale,
+                             flow_permutation='invconv',
+                             flow_coupling='noCoupling',
+                             LU_decomposed=LU_decomposed, opt=opt))
+                self.output_shapes.append(
+                    [-1, self.C, H, W])
+
+    def arch_squeeze(self, H, W):
+        self.C, H, W = self.C * 4, H // 2, W // 2
+        self.layers.append(flow.SqueezeLayer(factor=2))
+        self.output_shapes.append([-1, self.C, H, W])
+        return H, W
+
+    def get_flow_permutation(self, flow_permutation, opt):
+        flow_permutation = opt['network_G']['flow'].get('flow_permutation', 'invconv')
+        return flow_permutation
+
+    def get_affineInCh(self, opt_get):
+        affineInCh = opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'blocks']) or []
+        affineInCh = (len(affineInCh) + 1) * 64
+        return affineInCh
+
+    def check_image_shape(self):
+        assert self.C == 1 or self.C == 3, ("image_shape should be HWC, like (64, 64, 3)"
+                                            "self.C == 1 or self.C == 3")
+
+    def forward(self, gt=None, rrdbResults=None, z=None, epses=None, logdet=0., reverse=False, eps_std=None,
+                y_onehot=None):
+
+        if reverse:
+            epses_copy = [eps for eps in epses] if isinstance(epses, list) else epses
+
+            sr, logdet = self.decode(rrdbResults, z, eps_std, epses=epses_copy, logdet=logdet, y_onehot=y_onehot)
+            return sr, logdet
+        else:
+            assert gt is not None
+            assert rrdbResults is not None
+            z, logdet = self.encode(gt, rrdbResults, logdet=logdet, epses=epses, y_onehot=y_onehot)
+
+            return z, logdet
+
+    def encode(self, gt, rrdbResults, logdet=0.0, epses=None, y_onehot=None):
+        fl_fea = gt
+        reverse = False
+        level_conditionals = {}
+        bypasses = {}
+
+        L = opt_get(self.opt, ['network_G', 'flow', 'L'])
+
+        for level in range(1, L + 1):
+            bypasses[level] = torch.nn.functional.interpolate(gt, scale_factor=2 ** -level, mode='bilinear', align_corners=False)
+
+        for layer, shape in zip(self.layers, self.output_shapes):
+            size = shape[2]
+            level = int(np.log(160 / size) / np.log(2))
+
+            if level > 0 and level not in level_conditionals.keys():
+                level_conditionals[level] = rrdbResults[self.levelToName[level]]
+
+            level_conditionals[level] = rrdbResults[self.levelToName[level]]
+
+            if isinstance(layer, FlowStep):
+                fl_fea, logdet = layer(fl_fea, logdet, reverse=reverse, rrdbResults=level_conditionals[level])
+            elif isinstance(layer, Split2d):
+                fl_fea, logdet = self.forward_split2d(epses, fl_fea, layer, logdet, reverse, level_conditionals[level],
+                                                      y_onehot=y_onehot)
+            else:
+                fl_fea, logdet = layer(fl_fea, logdet, reverse=reverse)
+
+        z = fl_fea
+
+        if not isinstance(epses, list):
+            return z, logdet
+
+        epses.append(z)
+        return epses, logdet
+
+    def forward_preFlow(self, fl_fea, logdet, reverse):
+        if hasattr(self, 'preFlow'):
+            for l in self.preFlow:
+                fl_fea, logdet = l(fl_fea, logdet, reverse=reverse)
+        return fl_fea, logdet
+
+    def forward_split2d(self, epses, fl_fea, layer, logdet, reverse, rrdbResults, y_onehot=None):
+        ft = None if layer.position is None else rrdbResults[layer.position]
+        fl_fea, logdet, eps = layer(fl_fea, logdet, reverse=reverse, eps=epses, ft=ft, y_onehot=y_onehot)
+
+        if isinstance(epses, list):
+            epses.append(eps)
+        return fl_fea, logdet
+
+    def decode(self, rrdbResults, z, eps_std=None, epses=None, logdet=0.0, y_onehot=None):
+        z = epses.pop() if isinstance(epses, list) else z
+
+        fl_fea = z
+        # debug.imwrite("fl_fea", fl_fea)
+        bypasses = {}
+        level_conditionals = {}
+        if not opt_get(self.opt, ['network_G', 'flow', 'levelConditional', 'conditional']) == True:
+            for level in range(self.L + 1):
+                level_conditionals[level] = rrdbResults[self.levelToName[level]]
+
+        for layer, shape in zip(reversed(self.layers), reversed(self.output_shapes)):
+            size = shape[2]
+            level = int(np.log(160 / size) / np.log(2))
+            # size = fl_fea.shape[2]
+            # level = int(np.log(160 / size) / np.log(2))
+
+            if isinstance(layer, Split2d):
+                fl_fea, logdet = self.forward_split2d_reverse(eps_std, epses, fl_fea, layer,
+                                                              rrdbResults[self.levelToName[level]], logdet=logdet,
+                                                              y_onehot=y_onehot)
+            elif isinstance(layer, FlowStep):
+                fl_fea, logdet = layer(fl_fea, logdet=logdet, reverse=True, rrdbResults=level_conditionals[level])
+            else:
+                fl_fea, logdet = layer(fl_fea, logdet=logdet, reverse=True)
+
+        sr = fl_fea
+
+        assert sr.shape[1] == 3
+        return sr, logdet
+
+    def forward_split2d_reverse(self, eps_std, epses, fl_fea, layer, rrdbResults, logdet, y_onehot=None):
+        ft = None if layer.position is None else rrdbResults[layer.position]
+        fl_fea, logdet = layer(fl_fea, logdet=logdet, reverse=True,
+                               eps=epses.pop() if isinstance(epses, list) else None,
+                               eps_std=eps_std, ft=ft, y_onehot=y_onehot)
+        return fl_fea, logdet
+
+
+def get_position_name(H, scale):
+    downscale_factor = 160 // H
+    position_name = 'fea_up{}'.format(scale / downscale_factor)
+    return position_name

models/SRFlow/code/models/modules/Permutations.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import numpy as np
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+import sys
+sys.path.append('../../../..')
+from SRFlow.code.models.modules import thops
+
+
+class InvertibleConv1x1(nn.Module):
+    def __init__(self, num_channels, LU_decomposed=False):
+        super().__init__()
+        w_shape = [num_channels, num_channels]
+        w_init = np.linalg.qr(np.random.randn(*w_shape))[0].astype(np.float32)
+        self.register_parameter("weight", nn.Parameter(torch.Tensor(w_init)))
+        self.w_shape = w_shape
+        self.LU = LU_decomposed
+
+    def get_weight(self, input, reverse):
+        w_shape = self.w_shape
+        pixels = thops.pixels(input)
+        dlogdet = torch.slogdet(self.weight)[1] * pixels
+        if not reverse:
+            weight = self.weight.view(w_shape[0], w_shape[1], 1, 1)
+        else:
+            weight = torch.inverse(self.weight.double()).float() \
+                .view(w_shape[0], w_shape[1], 1, 1)
+        return weight, dlogdet
+    def forward(self, input, logdet=None, reverse=False):
+        """
+        log-det = log|abs(|W|)| * pixels
+        """
+        weight, dlogdet = self.get_weight(input, reverse)
+        if not reverse:
+            z = F.conv2d(input, weight)
+            if logdet is not None:
+                logdet = logdet + dlogdet
+            return z, logdet
+        else:
+            z = F.conv2d(input, weight)
+            if logdet is not None:
+                logdet = logdet - dlogdet
+            return z, logdet

models/SRFlow/code/models/modules/RRDBNet_arch.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import functools
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import sys
+sys.path.append('../..')
+import SRFlow.code.models.modules.module_util as mutil
+from SRFlow.code.utils.util import opt_get
+
+
+class ResidualDenseBlock_5C(nn.Module):
+    def __init__(self, nf=64, gc=32, bias=True):
+        super(ResidualDenseBlock_5C, self).__init__()
+        # gc: growth channel, i.e. intermediate channels
+        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
+        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
+        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    '''Residual in Residual Dense Block'''
+
+    def __init__(self, nf, gc=32):
+        super(RRDB, self).__init__()
+        self.RDB1 = ResidualDenseBlock_5C(nf, gc)
+        self.RDB2 = ResidualDenseBlock_5C(nf, gc)
+        self.RDB3 = ResidualDenseBlock_5C(nf, gc)
+
+    def forward(self, x):
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    def __init__(self, in_nc, out_nc, nf, nb, gc=32, scale=4, opt=None):
+        self.opt = opt
+        super(RRDBNet, self).__init__()
+        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
+        self.scale = scale
+
+        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+        self.RRDB_trunk = mutil.make_layer(RRDB_block_f, nb)
+        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        #### upsampling
+        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        if self.scale >= 8:
+            self.upconv3 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        if self.scale >= 16:
+            self.upconv4 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        if self.scale >= 32:
+            self.upconv5 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+
+        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x, get_steps=False):
+        fea = self.conv_first(x)
+
+        block_idxs = opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'blocks']) or []
+        block_results = {}
+
+        for idx, m in enumerate(self.RRDB_trunk.children()):
+            fea = m(fea)
+            for b in block_idxs:
+                if b == idx:
+                    block_results["block_{}".format(idx)] = fea
+
+        trunk = self.trunk_conv(fea)
+
+        last_lr_fea = fea + trunk
+
+        fea_up2 = self.upconv1(F.interpolate(last_lr_fea, scale_factor=2, mode='nearest'))
+        fea = self.lrelu(fea_up2)
+
+        fea_up4 = self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest'))
+        fea = self.lrelu(fea_up4)
+
+        fea_up8 = None
+        fea_up16 = None
+        fea_up32 = None
+
+        if self.scale >= 8:
+            fea_up8 = self.upconv3(F.interpolate(fea, scale_factor=2, mode='nearest'))
+            fea = self.lrelu(fea_up8)
+        if self.scale >= 16:
+            fea_up16 = self.upconv4(F.interpolate(fea, scale_factor=2, mode='nearest'))
+            fea = self.lrelu(fea_up16)
+        if self.scale >= 32:
+            fea_up32 = self.upconv5(F.interpolate(fea, scale_factor=2, mode='nearest'))
+            fea = self.lrelu(fea_up32)
+
+        out = self.conv_last(self.lrelu(self.HRconv(fea)))
+
+        results = {'last_lr_fea': last_lr_fea,
+                   'fea_up1': last_lr_fea,
+                   'fea_up2': fea_up2,
+                   'fea_up4': fea_up4,
+                   'fea_up8': fea_up8,
+                   'fea_up16': fea_up16,
+                   'fea_up32': fea_up32,
+                   'out': out}
+
+        fea_up0_en = opt_get(self.opt, ['network_G', 'flow', 'fea_up0']) or False
+        if fea_up0_en:
+            results['fea_up0'] = F.interpolate(last_lr_fea, scale_factor=1/2, mode='bilinear', align_corners=False, recompute_scale_factor=True)
+        fea_upn1_en = opt_get(self.opt, ['network_G', 'flow', 'fea_up-1']) or False
+        if fea_upn1_en:
+            results['fea_up-1'] = F.interpolate(last_lr_fea, scale_factor=1/4, mode='bilinear', align_corners=False, recompute_scale_factor=True)
+
+        if get_steps:
+            for k, v in block_results.items():
+                results[k] = v
+            return results
+        else:
+            return out

models/SRFlow/code/models/modules/SRFlowNet_arch.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import sys
+sys.path.append('../../../..')
+from SRFlow.code.models.modules.RRDBNet_arch import RRDBNet
+from SRFlow.code.models.modules.FlowUpsamplerNet import FlowUpsamplerNet
+import SRFlow.code.models.modules.thops as thops
+import SRFlow.code.models.modules.flow as flow
+from SRFlow.code.utils.util import opt_get
+
+
+class SRFlowNet(nn.Module):
+    def __init__(self, in_nc, out_nc, nf, nb, gc=32, scale=4, K=None, opt=None, step=None):
+        super(SRFlowNet, self).__init__()
+
+        self.opt = opt
+        self.quant = 255 if opt_get(opt, ['datasets', 'train', 'quant']) is \
+                            None else opt_get(opt, ['datasets', 'train', 'quant'])
+        self.RRDB = RRDBNet(in_nc, out_nc, nf, nb, gc, scale, opt)
+        hidden_channels = opt_get(opt, ['network_G', 'flow', 'hidden_channels'])
+        hidden_channels = hidden_channels or 64
+        self.RRDB_training = True  # Default is true
+
+        train_RRDB_delay = opt_get(self.opt, ['network_G', 'train_RRDB_delay'])
+        set_RRDB_to_train = False
+        if set_RRDB_to_train:
+            self.set_rrdb_training(True)
+
+        self.flowUpsamplerNet = \
+            FlowUpsamplerNet((160, 160, 3), hidden_channels, K,
+                             flow_coupling=opt['network_G']['flow']['coupling'], opt=opt)
+        self.i = 0
+
+    def set_rrdb_training(self, trainable):
+        if self.RRDB_training != trainable:
+            for p in self.RRDB.parameters():
+                p.requires_grad = trainable
+            self.RRDB_training = trainable
+            return True
+        return False
+
+    def forward(self, gt=None, lr=None, z=None, eps_std=None, reverse=False, epses=None, reverse_with_grad=False,
+                lr_enc=None,
+                add_gt_noise=False, step=None, y_label=None):
+        if not reverse:
+            return self.normal_flow(gt, lr, epses=epses, lr_enc=lr_enc, add_gt_noise=add_gt_noise, step=step,
+                                    y_onehot=y_label)
+        else:
+            # assert lr.shape[0] == 1
+            assert lr.shape[1] == 3
+            # assert lr.shape[2] == 20
+            # assert lr.shape[3] == 20
+            # assert z.shape[0] == 1
+            # assert z.shape[1] == 3 * 8 * 8
+            # assert z.shape[2] == 20
+            # assert z.shape[3] == 20
+            if reverse_with_grad:
+                return self.reverse_flow(lr, z, y_onehot=y_label, eps_std=eps_std, epses=epses, lr_enc=lr_enc,
+                                         add_gt_noise=add_gt_noise)
+            else:
+                with torch.no_grad():
+                    return self.reverse_flow(lr, z, y_onehot=y_label, eps_std=eps_std, epses=epses, lr_enc=lr_enc,
+                                             add_gt_noise=add_gt_noise)
+
+    def normal_flow(self, gt, lr, y_onehot=None, epses=None, lr_enc=None, add_gt_noise=True, step=None):
+        if lr_enc is None:
+            lr_enc = self.rrdbPreprocessing(lr)
+
+        logdet = torch.zeros_like(gt[:, 0, 0, 0])
+        pixels = thops.pixels(gt)
+
+        z = gt
+
+        if add_gt_noise:
+            # Setup
+            noiseQuant = opt_get(self.opt, ['network_G', 'flow', 'augmentation', 'noiseQuant'], True)
+            if noiseQuant:
+                z = z + ((torch.rand(z.shape, device=z.device) - 0.5) / self.quant)
+            logdet = logdet + float(-np.log(self.quant) * pixels)
+
+        # Encode
+        epses, logdet = self.flowUpsamplerNet(rrdbResults=lr_enc, gt=z, logdet=logdet, reverse=False, epses=epses,
+                                              y_onehot=y_onehot)
+
+        objective = logdet.clone()
+
+        if isinstance(epses, (list, tuple)):
+            z = epses[-1]
+        else:
+            z = epses
+
+        objective = objective + flow.GaussianDiag.logp(None, None, z)
+
+        nll = (-objective) / float(np.log(2.) * pixels)
+
+        if isinstance(epses, list):
+            return epses, nll, logdet
+        return z, nll, logdet
+
+    def rrdbPreprocessing(self, lr):
+        rrdbResults = self.RRDB(lr, get_steps=True)
+        block_idxs = opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'blocks']) or []
+        if len(block_idxs) > 0:
+            concat = torch.cat([rrdbResults["block_{}".format(idx)] for idx in block_idxs], dim=1)
+
+            if opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'concat']) or False:
+                keys = ['last_lr_fea', 'fea_up1', 'fea_up2', 'fea_up4']
+                if 'fea_up0' in rrdbResults.keys():
+                    keys.append('fea_up0')
+                if 'fea_up-1' in rrdbResults.keys():
+                    keys.append('fea_up-1')
+                if self.opt['scale'] >= 8:
+                    keys.append('fea_up8')
+                if self.opt['scale'] == 16:
+                    keys.append('fea_up16')
+                for k in keys:
+                    h = rrdbResults[k].shape[2]
+                    w = rrdbResults[k].shape[3]
+                    rrdbResults[k] = torch.cat([rrdbResults[k], F.interpolate(concat, (h, w))], dim=1)
+        return rrdbResults
+
+    def get_score(self, disc_loss_sigma, z):
+        score_real = 0.5 * (1 - 1 / (disc_loss_sigma ** 2)) * thops.sum(z ** 2, dim=[1, 2, 3]) - \
+                     z.shape[1] * z.shape[2] * z.shape[3] * math.log(disc_loss_sigma)
+        return -score_real
+
+    def reverse_flow(self, lr, z, y_onehot, eps_std, epses=None, lr_enc=None, add_gt_noise=True):
+        logdet = torch.zeros_like(lr[:, 0, 0, 0])
+        pixels = thops.pixels(lr) * self.opt['scale'] ** 2
+
+        if add_gt_noise:
+            logdet = logdet - float(-np.log(self.quant) * pixels)
+
+        if lr_enc is None:
+            lr_enc = self.rrdbPreprocessing(lr)
+
+        x, logdet = self.flowUpsamplerNet(rrdbResults=lr_enc, z=z, eps_std=eps_std, reverse=True, epses=epses,
+                                          logdet=logdet)
+
+        return x, logdet

models/SRFlow/code/models/modules/Split.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+from torch import nn as nn
+import sys
+sys.path.append('../../../..')
+from SRFlow.code.models.modules import thops
+from SRFlow.code.models.modules.FlowStep import FlowStep
+from SRFlow.code.models.modules.flow import Conv2dZeros, GaussianDiag
+from SRFlow.code.utils.util import opt_get
+
+
+class Split2d(nn.Module):
+    def __init__(self, num_channels, logs_eps=0, cond_channels=0, position=None, consume_ratio=0.5, opt=None):
+        super().__init__()
+
+        self.num_channels_consume = int(round(num_channels * consume_ratio))
+        self.num_channels_pass = num_channels - self.num_channels_consume
+
+        self.conv = Conv2dZeros(in_channels=self.num_channels_pass + cond_channels,
+                                out_channels=self.num_channels_consume * 2)
+        self.logs_eps = logs_eps
+        self.position = position
+        self.opt = opt
+
+    def split2d_prior(self, z, ft):
+        if ft is not None:
+            z = torch.cat([z, ft], dim=1)
+        h = self.conv(z)
+        return thops.split_feature(h, "cross")
+
+    def exp_eps(self, logs):
+        return torch.exp(logs) + self.logs_eps
+
+    def forward(self, input, logdet=0., reverse=False, eps_std=None, eps=None, ft=None, y_onehot=None):
+        if not reverse:
+            # self.input = input
+            z1, z2 = self.split_ratio(input)
+            mean, logs = self.split2d_prior(z1, ft)
+            
+            eps = (z2 - mean) / self.exp_eps(logs)
+
+            logdet = logdet + self.get_logdet(logs, mean, z2)
+
+            # print(logs.shape, mean.shape, z2.shape)
+            # self.eps = eps
+            # print('split, enc eps:', eps)
+            return z1, logdet, eps
+        else:
+            z1 = input
+            mean, logs = self.split2d_prior(z1, ft)
+
+            if eps is None:
+                #print("WARNING: eps is None, generating eps untested functionality!")
+                eps = GaussianDiag.sample_eps(mean.shape, eps_std)
+
+            eps = eps.to(mean.device)
+            z2 = mean + self.exp_eps(logs) * eps
+
+            z = thops.cat_feature(z1, z2)
+            logdet = logdet - self.get_logdet(logs, mean, z2)
+
+            return z, logdet
+            # return z, logdet, eps
+
+    def get_logdet(self, logs, mean, z2):
+        logdet_diff = GaussianDiag.logp(mean, logs, z2)
+        # print("Split2D: logdet diff", logdet_diff.item())
+        return logdet_diff
+
+    def split_ratio(self, input):
+        z1, z2 = input[:, :self.num_channels_pass, ...], input[:, self.num_channels_pass:, ...]
+        return z1, z2

models/SRFlow/code/models/modules/flow.py

@@ -0,0 +1,167 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import sys
+sys.path.append('../../../..')
+from SRFlow.code.models.modules.FlowActNorms import ActNorm2d
+from . import thops
+
+
+class Conv2d(nn.Conv2d):
+    pad_dict = {
+        "same": lambda kernel, stride: [((k - 1) * s + 1) // 2 for k, s in zip(kernel, stride)],
+        "valid": lambda kernel, stride: [0 for _ in kernel]
+    }
+
+    @staticmethod
+    def get_padding(padding, kernel_size, stride):
+        # make paddding
+        if isinstance(padding, str):
+            if isinstance(kernel_size, int):
+                kernel_size = [kernel_size, kernel_size]
+            if isinstance(stride, int):
+                stride = [stride, stride]
+            padding = padding.lower()
+            try:
+                padding = Conv2d.pad_dict[padding](kernel_size, stride)
+            except KeyError:
+                raise ValueError("{} is not supported".format(padding))
+        return padding
+
+    def __init__(self, in_channels, out_channels,
+                 kernel_size=[3, 3], stride=[1, 1],
+                 padding="same", do_actnorm=True, weight_std=0.05):
+        padding = Conv2d.get_padding(padding, kernel_size, stride)
+        super().__init__(in_channels, out_channels, kernel_size, stride,
+                         padding, bias=(not do_actnorm))
+        # init weight with std
+        self.weight.data.normal_(mean=0.0, std=weight_std)
+        if not do_actnorm:
+            self.bias.data.zero_()
+        else:
+            self.actnorm = ActNorm2d(out_channels)
+        self.do_actnorm = do_actnorm
+
+    def forward(self, input):
+        x = super().forward(input)
+        if self.do_actnorm:
+            x, _ = self.actnorm(x)
+        return x
+
+
+class Conv2dZeros(nn.Conv2d):
+    def __init__(self, in_channels, out_channels,
+                 kernel_size=[3, 3], stride=[1, 1],
+                 padding="same", logscale_factor=3):
+        padding = Conv2d.get_padding(padding, kernel_size, stride)
+        super().__init__(in_channels, out_channels, kernel_size, stride, padding)
+        # logscale_factor
+        self.logscale_factor = logscale_factor
+        self.register_parameter("logs", nn.Parameter(torch.zeros(out_channels, 1, 1)))
+        # init
+        self.weight.data.zero_()
+        self.bias.data.zero_()
+
+    def forward(self, input):
+        output = super().forward(input)
+        return output * torch.exp(self.logs * self.logscale_factor)
+
+
+class GaussianDiag:
+    Log2PI = float(np.log(2 * np.pi))
+
+    @staticmethod
+    def likelihood(mean, logs, x):
+        """
+        lnL = -1/2 * { ln|Var| + ((X - Mu)^T)(Var^-1)(X - Mu) + kln(2*PI) }
+              k = 1 (Independent)
+              Var = logs ** 2
+        """
+        if mean is None and logs is None:
+            return -0.5 * (x ** 2 + GaussianDiag.Log2PI)
+        else:
+            return -0.5 * (logs * 2. + ((x - mean) ** 2) / torch.exp(logs * 2.) + GaussianDiag.Log2PI)
+
+    @staticmethod
+    def logp(mean, logs, x):
+        likelihood = GaussianDiag.likelihood(mean, logs, x)
+        return thops.sum(likelihood, dim=[1, 2, 3])
+
+    @staticmethod
+    def sample(mean, logs, eps_std=None):
+        eps_std = eps_std or 1
+        eps = torch.normal(mean=torch.zeros_like(mean),
+                           std=torch.ones_like(logs) * eps_std)
+        return mean + torch.exp(logs) * eps
+
+    @staticmethod
+    def sample_eps(shape, eps_std, seed=None):
+        if seed is not None:
+            torch.manual_seed(seed)
+        eps = torch.normal(mean=torch.zeros(shape),
+                           std=torch.ones(shape) * eps_std)
+        return eps
+
+
+def squeeze2d(input, factor=2):
+    assert factor >= 1 and isinstance(factor, int)
+    if factor == 1:
+        return input
+    size = input.size()
+    B = size[0]
+    C = size[1]
+    H = size[2]
+    W = size[3]
+    assert H % factor == 0 and W % factor == 0, "{}".format((H, W, factor))
+    x = input.view(B, C, H // factor, factor, W // factor, factor)
+    x = x.permute(0, 1, 3, 5, 2, 4).contiguous()
+    x = x.view(B, C * factor * factor, H // factor, W // factor)
+    return x
+
+
+def unsqueeze2d(input, factor=2):
+    assert factor >= 1 and isinstance(factor, int)
+    factor2 = factor ** 2
+    if factor == 1:
+        return input
+    size = input.size()
+    B = size[0]
+    C = size[1]
+    H = size[2]
+    W = size[3]
+    assert C % (factor2) == 0, "{}".format(C)
+    x = input.view(B, C // factor2, factor, factor, H, W)
+    x = x.permute(0, 1, 4, 2, 5, 3).contiguous()
+    x = x.view(B, C // (factor2), H * factor, W * factor)
+    return x
+
+
+class SqueezeLayer(nn.Module):
+    def __init__(self, factor):
+        super().__init__()
+        self.factor = factor
+
+    def forward(self, input, logdet=None, reverse=False):
+        if not reverse:
+            output = squeeze2d(input, self.factor)  # Squeeze in forward
+            return output, logdet
+        else:
+            output = unsqueeze2d(input, self.factor)
+            return output, logdet

models/SRFlow/code/models/modules/glow_arch.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/chaiyujin/glow-pytorch/blob/master/LICENSE
+
+import torch.nn as nn
+
+
+def f_conv2d_bias(in_channels, out_channels):
+    def padding_same(kernel, stride):
+        return [((k - 1) * s + 1) // 2 for k, s in zip(kernel, stride)]
+
+    padding = padding_same([3, 3], [1, 1])
+    assert padding == [1, 1], padding
+    return nn.Sequential(
+        nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=[3, 3], stride=1, padding=1,
+                  bias=True))