removed v3 gan

core/leras/models/PatchDiscriminator.py

@@ -76,134 +76,9 @@ class PatchDiscriminator(nn.ModelBase):
 
 nn.PatchDiscriminator = PatchDiscriminator
 
-
 class UNetPatchDiscriminator(nn.ModelBase):
     """
     Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks"
-    Based on iperov commit 11add4cd4f5a61df26a8659f4cc5c8d9467bf5f8 from Jan 3, 2021 with added fp16 option
-    """
-    def calc_receptive_field_size(self, layers):
-        """
-        result the same as https://fomoro.com/research/article/receptive-field-calculatorindex.html
-        """
-        rf = 0
-        ts = 1
-        for i, (k, s) in enumerate(layers):
-            if i == 0:
-                rf = k
-            else:
-                rf += (k-1)*ts
-            ts *= s
-        return rf
-
-    def find_archi(self, target_patch_size, max_layers=9):
-        """
-        Find the best configuration of layers using only 3x3 convs for target patch size
-        """
-        s = {}
-        for layers_count in range(1,max_layers+1):
-            val = 1 << (layers_count-1)
-            while True:
-                val -= 1
-
-                layers = []
-                sum_st = 0
-                layers.append ( [3, 2])
-                sum_st += 2
-                for i in range(layers_count-1):
-                    st = 1 + (1 if val & (1 << i) !=0 else 0 )
-                    layers.append ( [3, st ])
-                    sum_st += st
-
-                rf = self.calc_receptive_field_size(layers)
-
-                s_rf = s.get(rf, None)
-                if s_rf is None:
-                    s[rf] = (layers_count, sum_st, layers)
-                else:
-                    if layers_count < s_rf[0] or \
-                    ( layers_count == s_rf[0] and sum_st > s_rf[1] ):
-                        s[rf] = (layers_count, sum_st, layers)
-
-                if val == 0:
-                    break
-
-        x = sorted(list(s.keys()))
-        q=x[np.abs(np.array(x)-target_patch_size).argmin()]
-        return s[q][2]
-
-    def on_build(self, patch_size, in_ch, base_ch = 16, use_fp16 = False):
-        self.use_fp16 = use_fp16
-        conv_dtype = tf.float16 if use_fp16 else tf.float32 
-    
-        class ResidualBlock(nn.ModelBase):
-            def on_build(self, ch, kernel_size=3 ):
-                self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
-                self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
-
-            def forward(self, inp):
-                x = self.conv1(inp)
-                x = tf.nn.leaky_relu(x, 0.2)
-                x = self.conv2(x)
-                x = tf.nn.leaky_relu(inp + x, 0.2)
-                return x
-
-        prev_ch = in_ch
-        self.convs = []
-        self.res1 = []
-        self.res2 = []
-        self.upconvs = []
-        self.upres1 = []
-        self.upres2 = []
-        layers = self.find_archi(patch_size)
-        
-        level_chs = { i-1:v for i,v in enumerate([ min( base_ch * (2**i), 512 ) for i in range(len(layers)+1)]) }
-
-        self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID', dtype=conv_dtype)
-
-        for i, (kernel_size, strides) in enumerate(layers):
-            self.convs.append ( nn.Conv2D( level_chs[i-1], level_chs[i], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) )
-
-            self.res1.append ( ResidualBlock(level_chs[i]) )
-            self.res2.append ( ResidualBlock(level_chs[i]) )
-            
-            self.upconvs.insert (0, nn.Conv2DTranspose( level_chs[i]*(2 if i != len(layers)-1 else 1), level_chs[i-1], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) )
-
-            self.upres1.insert (0, ResidualBlock(level_chs[i-1]*2) )
-            self.upres2.insert (0, ResidualBlock(level_chs[i-1]*2) )
-            
-        self.out_conv = nn.Conv2D( level_chs[-1]*2, 1, kernel_size=1, padding='VALID', dtype=conv_dtype)
-
-        self.center_out  =  nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID', dtype=conv_dtype)
-        self.center_conv =  nn.Conv2D( level_chs[len(layers)-1], level_chs[len(layers)-1], kernel_size=1, padding='VALID', dtype=conv_dtype)
-
-
-    def forward(self, x):
-        x = tf.nn.leaky_relu( self.in_conv(x), 0.2 )
-
-        encs = []
-        for conv, res1,res2 in zip(self.convs, self.res1, self.res2):
-            encs.insert(0, x)
-            x = tf.nn.leaky_relu( conv(x), 0.2 )
-            x = res1(x)
-            x = res2(x)
-            
-        center_out, x = self.center_out(x), tf.nn.leaky_relu( self.center_conv(x), 0.2 )
-
-        for i, (upconv, enc, upres1, upres2 ) in enumerate(zip(self.upconvs, encs, self.upres1, self.upres2)):
-            x = tf.nn.leaky_relu( upconv(x), 0.2 )
-            x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis)
-            x = upres1(x)
-            x = upres2(x)
-
-        return center_out, self.out_conv(x)
-
-nn.UNetPatchDiscriminator = UNetPatchDiscriminator
-
-class UNetPatchDiscriminatorV2(nn.ModelBase):
-    """
-    Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks"
-    Based on iperov commit 35877dbfd724c22040f421e93c1adbb7142e5b5d from Jul 14, 2021
     """
     def calc_receptive_field_size(self, layers):
         """
@@ -316,4 +191,4 @@ class UNetPatchDiscriminatorV2(nn.ModelBase):
 
         return center_out, x
 
-nn.UNetPatchDiscriminatorV2 = UNetPatchDiscriminatorV2
+nn.UNetPatchDiscriminator = UNetPatchDiscriminator

models/Model_AMP/Model.py

@@ -132,7 +132,6 @@ class AMPModel(ModelBase):
                 self.options['lr_dropout']  = io.input_str (f"Use learning rate dropout", default_lr_dropout, ['n','y','cpu'], help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations. Enabled it before `disable random warp` and before GAN. \nn - disabled.\ny - enabled\ncpu - enabled on CPU. This allows not to use extra VRAM, sacrificing 20% time of iteration.")
 
         default_gan_power          = self.options['gan_power']          = self.load_or_def_option('gan_power', 0.0)
-        default_gan_version        = self.options['gan_version']        = self.load_or_def_option('gan_version', 2)        
         default_gan_patch_size     = self.options['gan_patch_size']     = self.load_or_def_option('gan_patch_size', self.options['resolution'] // 8)
         default_gan_dims           = self.options['gan_dims']           = self.load_or_def_option('gan_dims', 16)
         default_gan_smoothing      = self.options['gan_smoothing']      = self.load_or_def_option('gan_smoothing', 0.1)
@@ -154,9 +153,7 @@ class AMPModel(ModelBase):
 
 
                 if self.options['gan_power'] != 0.0:
-                    self.options['gan_version'] = np.clip (io.input_int("GAN version", default_gan_version, add_info="2 or 3", help_message="Choose GAN version (v2: 7/16/2020, v3: 1/3/2021):"), 2, 3)
 
-                    if self.options['gan_version'] == 3:
                     gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 )
                     self.options['gan_patch_size'] = gan_patch_size
 
@@ -381,10 +378,6 @@ class AMPModel(ModelBase):
 
             if self.is_training:
                 if gan_power != 0:
-                    if self.options['gan_version'] == 2:
-                        self.GAN = nn.UNetPatchDiscriminatorV2(patch_size=resolution//16, in_ch=input_ch, name="D_src", use_fp16=self.options['use_fp16'])
-                        self.model_filename_list += [ [self.GAN, 'D_src_v2.npy'] ]
-                    else:
                     self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=self.options['use_fp16'], name="D_src")
                     self.model_filename_list += [ [self.GAN, 'GAN.npy'] ]
             
@@ -411,11 +404,6 @@ class AMPModel(ModelBase):
                 self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
 
                 if gan_power != 0:
-                    if self.options['gan_version'] == 2:
-                        self.GAN_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='D_src_dst_opt')
-                        self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
-                        self.model_filename_list += [ (self.GAN_opt, 'D_src_v2_opt.npy') ]
-                    else:
                     self.GAN_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt')
                     self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
                     self.model_filename_list += [ (self.GAN_opt, 'GAN_opt.npy') ]

models/Model_SAEHD/Model.py

@@ -165,7 +165,6 @@ class SAEHDModel(ModelBase):
                 self.options['blur_out_mask'] = io.input_bool ("Blur out mask", default_blur_out_mask, help_message='Blurs nearby area outside of applied face mask of training samples. The result is the background near the face is smoothed and less noticeable on swapped face. The exact xseg mask in src and dst faceset is required.')
 
         default_gan_power          = self.options['gan_power']          = self.load_or_def_option('gan_power', 0.0)
-        default_gan_version        = self.options['gan_version']        = self.load_or_def_option('gan_version', 2)
         default_gan_patch_size     = self.options['gan_patch_size']     = self.load_or_def_option('gan_patch_size', self.options['resolution'] // 8)
         default_gan_dims           = self.options['gan_dims']           = self.load_or_def_option('gan_dims', 16)
         default_gan_smoothing      = self.options['gan_smoothing']      = self.load_or_def_option('gan_smoothing', 0.1)
@@ -194,9 +193,7 @@ class SAEHDModel(ModelBase):
                 self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 10.0", help_message="Train the network in Generative Adversarial manner. Forces the neural network to learn small details of the face. Enable it only when the face is trained enough and don't disable. Typical value is 0.1"), 0.0, 10.0 )
 
                 if self.options['gan_power'] != 0.0:
-                    self.options['gan_version'] = np.clip (io.input_int("GAN version", default_gan_version, add_info="2 or 3", help_message="Choose GAN version (v2: 7/16/2020, v3: 1/3/2021):"), 2, 3)
 
-                    if self.options['gan_version'] == 3:
                     gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 )
                     self.options['gan_patch_size'] = gan_patch_size
 
@@ -355,10 +352,6 @@ class SAEHDModel(ModelBase):
 
             if self.is_training:
                 if gan_power != 0:
-                    if self.options['gan_version'] == 2:
-                        self.D_src = nn.UNetPatchDiscriminatorV2(patch_size=resolution//16, in_ch=input_ch, name="D_src", use_fp16=self.options['use_fp16'])
-                        self.model_filename_list += [ [self.D_src, 'D_src_v2.npy'] ]
-                    else:
                     self.D_src = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=self.options['use_fp16'], name="D_src")
                     self.model_filename_list += [ [self.D_src, 'GAN.npy'] ]
 
@@ -400,15 +393,11 @@ class SAEHDModel(ModelBase):
                     self.model_filename_list += [ (self.D_code_opt, 'D_code_opt.npy') ]
 
                 if gan_power != 0:
-                    if self.options['gan_version'] == 2:
-                        self.D_src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='D_src_dst_opt')
-                        self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
-                        self.model_filename_list += [ (self.D_src_dst_opt, 'D_src_v2_opt.npy') ]
-                    else:
                     self.D_src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt')
                     self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
                     self.model_filename_list += [ (self.D_src_dst_opt, 'GAN_opt.npy') ]
 
+
         if self.is_training:
             # Adjust batch size for multiple GPU
             gpu_count = max(1, len(devices) )