Update app.py

app.py

@@ -7,8 +7,8 @@ from PIL import Image
 vgg = tf.keras.applications.VGG19(include_top=False, weights="imagenet")
 vgg.trainable = False
 
-STYLE_LAYERS = [...]  # same layers as in your notebook
-CONTENT_LAYER = [...] # e.g. [('block5_conv4', 1)]
+STYLE_LAYERS = [...]  # your style layers
+CONTENT_LAYER = [...] # your content layer(s)
 
 
 def gram_matrix(A):
@@ -18,63 +18,52 @@ def gram_matrix(A):
     A = tf.convert_to_tensor(A)
 
     if len(A.shape) == 4:
-        # (batch, H, W, C)
         batch, H, W, C = A.shape
-        A = tf.reshape(A, (batch, H*W, C))
-        gram = tf.matmul(A, A, transpose_a=True) / tf.cast(H*W, tf.float32)
-
+        A = tf.reshape(A, (batch, H * W, C))
+        gram = tf.matmul(A, A, transpose_a=True) / tf.cast(H * W, tf.float32)
     elif len(A.shape) == 3:
-        # (batch, N, C) already flattened
         batch, N, C = A.shape
         gram = tf.matmul(A, A, transpose_a=True) / tf.cast(N, tf.float32)
-
     elif len(A.shape) == 2:
-        # (N, C) no batch
         N, C = A.shape
-        A = tf.expand_dims(A, 0)  # add batch dim
+        A = tf.expand_dims(A, 0)  
         gram = tf.matmul(A, A, transpose_a=True) / tf.cast(N, tf.float32)
-
     else:
         raise ValueError(f"Unexpected tensor rank for gram_matrix: {A.shape}")
-
     return gram
+
+
 def compute_content_cost(a_C, a_G):
-    """Content cost between content and generated image features."""
     return tf.reduce_mean(tf.square(a_C - a_G))
 
+
 def compute_style_cost(a_S, a_G):
-    """Style cost using Gram matrices."""
     J_style = 0
     for s, g in zip(a_S, a_G):
         J_style += tf.reduce_mean(tf.square(gram_matrix(s) - gram_matrix(g)))
     return J_style / len(a_S)
 
+
 def total_cost(J_content, J_style, alpha=10, beta=40):
-    """Weighted sum of content + style."""
     return alpha * J_content + beta * J_style
 
-def clip_0_1(img):
-    """Keep pixel values in [0,1]."""
-    return tf.clip_by_value(img, 0.0, 1.0)
 
 def preprocess(img):
     img = Image.fromarray(img).resize((256, 256))
     arr = np.expand_dims(np.array(img) / 255.0, axis=0).astype(np.float32)
     return tf.convert_to_tensor(arr)
 
-def style_transfer(content, style, steps=1000):
+
+def style_transfer(content, style, steps):
     content_tensor = preprocess(content)
     style_tensor = preprocess(style)
 
-    # Extract features
     a_C = vgg(content_tensor)
     a_S = vgg(style_tensor)
 
-    # Init generated image
     generated_image = tf.Variable(content_tensor)
-
-    # Optimize
     opt = tf.keras.optimizers.Adam(learning_rate=0.01)
+
     for i in range(steps):
         with tf.GradientTape() as tape:
             a_G = vgg(generated_image)
@@ -89,11 +78,15 @@ def style_transfer(content, style, steps=1000):
     out_img = (generated_image[0].numpy() * 255).astype("uint8")
     return Image.fromarray(out_img)
 
+
 demo = gr.Interface(
     fn=style_transfer,
-    inputs=[gr.Image(type="numpy", label="Content Image"), 
-            gr.Image(type="numpy", label="Style Image")],
+    inputs=[
+        gr.Image(type="numpy", label="Content Image"), 
+        gr.Image(type="numpy", label="Style Image"),
+        gr.Slider(100, 2000, value=1000, step=100, label="Number of Iterations")
+    ],
     outputs=gr.Image(type="pil", label="Stylized Image"),
 )
 
-demo.launch()
+demo.launch()