Update app.py

app.py

@@ -1,9 +1,10 @@
+%%writefile st.py
 import streamlit as st
 import tensorflow as tf
 from tensorflow_addons.layers import InstanceNormalization
 import numpy as np
 from PIL import Image
-from huggingface_hub import from_pretrained_keras
+from huggingface_hub import hf_hub_download
 import requests
 from io import BytesIO
 
@@ -56,14 +57,54 @@ st.markdown(f"<style>{combined_css}</style>", unsafe_allow_html=True)
 st.markdown('<div class="title"><span class="colorful-text">Photo</span> <span class="black-white-text">to Art</span></div>', unsafe_allow_html=True)
 st.markdown('<div class="custom-text">Convert Photos to Art using CycleGAN</div>', unsafe_allow_html=True)
 
-# Define your Hugging Face repository details
-username = "Hammad712"  # Replace with your Hugging Face username
-repo_name = "CycleGAN-Model"
-model_id = f"{username}/{repo_name}"
-
+# Define CycleGAN model architecture
+class CycleGAN(tf.keras.Model):
+    def __init__(self):
+        super(CycleGAN, self).__init__()
+        self.generatorP = self.build_generator()  # Generates Photo from Style
+        self.generatorS = self.build_generator()  # Generates Style from Photo
+
+    def build_generator(self):
+        def ConvBlock(filters, kernel, strides, x):
+            x = tf.keras.layers.Conv2D(filters, kernel, strides, 'same', kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), use_bias=False)(x)
+            if filters == 3:
+                return tf.keras.layers.Activation('tanh')(x)
+            x = InstanceNormalization(gamma_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02))(x)
+            return tf.keras.layers.ReLU()(x)
+
+        def ResBlock(filters, inputs):
+            x = tf.keras.layers.Conv2D(filters, 3, 1, padding='same', kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), use_bias=False)(inputs)
+            x = InstanceNormalization(gamma_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02))(x)
+            x = tf.keras.layers.ReLU()(x)
+            x = tf.keras.layers.Conv2D(filters, 3, 1, padding='same', kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), use_bias=False)(x)
+            x = InstanceNormalization(gamma_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02))(x)
+            return tf.keras.layers.Add()([x, inputs])
+
+        def TransBlock(filters, x):
+            x = tf.keras.layers.Conv2DTranspose(filters, 3, 2, 'same', kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), use_bias=False)(x)
+            x = InstanceNormalization(gamma_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02))(x)
+            return tf.keras.layers.ReLU()(x)
+
+        input = tf.keras.Input(shape=(256, 256, 3))
+        x = ConvBlock(64, 7, 1, input)
+        x = ConvBlock(128, 3, 2, x)
+        x = ConvBlock(256, 3, 2, x)
+        for _ in range(9):
+            x = ResBlock(256, x)
+        x = TransBlock(128, x)
+        x = TransBlock(64, x)
+        out = ConvBlock(3, 7, 1, x)
+        return tf.keras.Model(inputs=input, outputs=out)
+
+# Load the model weights from Hugging Face
 try:
-    # Load the model
-    model = from_pretrained_keras(model_id)
+    repo_id = "Hammad712/CycleGAN-Model"
+    generatorS_path = hf_hub_download(repo_id=repo_id, filename="generatorS.h5")
+    generatorP_path = hf_hub_download(repo_id=repo_id, filename="generatorP.h5")
+
+    model = CycleGAN()
+    model.generatorS.load_weights(generatorS_path)
+    model.generatorP.load_weights(generatorP_path)
     st.success("Model loaded successfully!")
 except Exception as e:
     st.error(f"An error occurred while loading the model: {e}")