Update app.py

app.py

@@ -23,17 +23,14 @@ num_content_layers = len(content_layers)
 num_style_layers = len(style_layers)
 
 # Model for extracting style feature
-style_models = [
-    keras.Model(inputs=model.input, outputs=model.get_layer(layer).output) for layer in style_layers
-]
+style_model = keras.Model(inputs=model.input, outputs=[model.get_layer(name).output for name in style_layers], name='Style_model')
 
 # Model for extracting content feature
 content_model = keras.Model(inputs=model.input, outputs=model.get_layer(content_layers).output, name='Content_model')
 
-# Style models summary
-for m in style_models:
-    print(f'Style model: {m.name}')
-    print(m.summary())
+# Style model summary
+print(f'Style model: {style_model.name}')
+style_model.summary()
 
 # Content model summary
 print(f'Content model: {content_model.name}')
@@ -54,11 +51,13 @@ def gram_matrix(image):
 
 # Style loss
 def style_loss(gen, style_img):
-    # Compute Gram matrix of generated image and style image
-    gen_gram = gram_matrix(gen)
-    style_gram = gram_matrix(style_img)
-    # Compute style loss
-    return tf.reduce_sum(tf.square(gen_gram-style_gram))/4
+    total_loss = 0.
+    for i in range(len(gen)):
+        gen_gram = gram_matrix(gen[i])
+        style_gram = gram_matrix(style_img[i])
+        # Compute style loss
+        total_loss += tf.reduce_sum(tf.square(gen_gram-style_gram))/4
+    return total_loss
 
 # Training parameter
 lr = 10.
@@ -76,11 +75,9 @@ def train_step(image, process_content_image, process_style_image):
         content_feature  = content_model(process_content_image)
         loss_content = content_loss(gen_content_feature, content_feature)
         # Calculate style loss
-        loss_style = 0
-        for style_model in style_models:
-            gen_style_feature = style_model(image)
-            style_feature = style_model(process_style_image)
-            loss_style += style_loss(gen_style_feature, style_feature)
+        gen_style_feature = style_model(image)
+        style_feature = style_model(process_style_image)
+        loss_style = style_loss(gen_style_feature, style_feature)
         # Calculate total loss
         total_loss = alpha*loss_content + beta*loss_style/style_weight
         # Calculate gradient