Upload tf_models.py with huggingface_hub

tf_models.py

@@ -0,0 +1,76 @@
+import tensorflow as tf
+from tensorflow.keras import layers
+
+def downsample(filters, size, apply_instancenorm=True):
+    initializer = tf.random_normal_initializer(0., 0.02)
+    result = tf.keras.Sequential()
+    result.add(layers.Conv2D(filters, size, strides=2, padding='same',
+                             kernel_initializer=initializer, use_bias=False))
+    if apply_instancenorm:
+        result.add(tf.keras.layers.GroupNormalization(groups=-1))
+    result.add(layers.LeakyReLU())
+    return result
+
+def upsample(filters, size, apply_dropout=False):
+    initializer = tf.random_normal_initializer(0., 0.02)
+    result = tf.keras.Sequential()
+    result.add(layers.Conv2DTranspose(filters, size, strides=2, padding='same',
+                                      kernel_initializer=initializer, use_bias=False))
+    result.add(tf.keras.layers.GroupNormalization(groups=-1))
+    if apply_dropout:
+        result.add(layers.Dropout(0.5))
+    result.add(layers.ReLU())
+    return result
+
+def resnet_block(filters, size=3):
+    initializer = tf.random_normal_initializer(0., 0.02)
+    result = tf.keras.Sequential()
+    result.add(layers.Conv2D(filters, size, padding='same', kernel_initializer=initializer, use_bias=False))
+    result.add(tf.keras.layers.GroupNormalization(groups=-1))
+    result.add(layers.ReLU())
+    result.add(layers.Conv2D(filters, size, padding='same', kernel_initializer=initializer, use_bias=False))
+    result.add(tf.keras.layers.GroupNormalization(groups=-1))
+    return result
+
+def Generator(output_channels=3, num_resnet=9):
+    inputs = layers.Input(shape=[256, 256, 3])
+    
+    # Downsampling
+    x = layers.Conv2D(64, 7, padding='same', kernel_initializer=tf.random_normal_initializer(0., 0.02), use_bias=False)(inputs)
+    x = tf.keras.layers.GroupNormalization(groups=-1)(x)
+    x = layers.ReLU()(x)
+    
+    x = downsample(128, 3)(x)
+    x = downsample(256, 3)(x)
+    
+    # Residual blocks
+    for _ in range(num_resnet):
+        res = resnet_block(256)(x)
+        x = layers.Add()([x, res])
+        
+    # Upsampling
+    x = upsample(128, 3)(x)
+    x = upsample(64, 3)(x)
+    
+    last = layers.Conv2D(output_channels, 7, padding='same', activation='tanh', 
+                          kernel_initializer=tf.random_normal_initializer(0., 0.02))(x)
+    
+    return tf.keras.Model(inputs=inputs, outputs=last)
+
+def Discriminator():
+    initializer = tf.random_normal_initializer(0., 0.02)
+    inputs = layers.Input(shape=[256, 256, 3])
+    
+    down1 = downsample(64, 4, False)(inputs) # (bs, 128, 128, 64)
+    down2 = downsample(128, 4)(down1)       # (bs, 64, 64, 128)
+    down3 = downsample(256, 4)(down2)       # (bs, 32, 32, 256)
+    
+    zero_pad1 = layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)
+    conv = layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1)
+    norm1 = tf.keras.layers.GroupNormalization(groups=-1)(conv)
+    leaky_relu = layers.LeakyReLU()(norm1)
+    
+    zero_pad2 = layers.ZeroPadding2D()(leaky_relu)
+    last = layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer)(zero_pad2)
+    
+    return tf.keras.Model(inputs=inputs, outputs=last)