Update app.py

app.py

@@ -33,12 +33,14 @@ def handle_image_size_selection(img_size, content_img):
                 None)
 
 # Define the function to process images
+@spaces.GPU
 def process_images(content_img, style_img, epochs, steps_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height):
     print("Start processing")
-    output_img = main(content_img, style_img, epochs, steps_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height)
+    with device("/cuda:0"):
+        output_img = main(content_img, style_img, epochs, steps_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height)
     return output_img
 
-@spaces.GPU
+# @spaces.GPU
 def create_app():
     with gr.Blocks() as demo:
         aspect_ratio = gr.State(None)