Update app.py

app.py

@@ -4,6 +4,7 @@ import warnings
 import os
 import shutil
 import subprocess
+import spaces
 import torch
 import numpy as np
 from diffusers import FluxControlNetModel
@@ -95,6 +96,8 @@ logger.info("Pipeline loaded successfully.")
 MAX_SEED = 1000000
 MAX_PIXEL_BUDGET = 1024 * 1024
 
+
+@spaces.GPU
 def process_input(input_image, upscale_factor):
     w, h = input_image.size
     aspect_ratio = w / h
@@ -118,6 +121,7 @@ def process_input(input_image, upscale_factor):
 
     return input_image.resize((w, h)), was_resized
 
+@spaces.GPU
 def run_inference(process_id, input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale):
     logger.info("Processing inference for process_id: %s", process_id)
     input_image, was_resized = process_input(input_image, upscale_factor)
@@ -157,6 +161,7 @@ def run_inference(process_id, input_image, upscale_factor, seed, num_inference_s
     logger.info("Inference completed for process_id: %s", process_id)
 
 @app.route('/infer', methods=['POST'])
+@spaces.GPU
 def infer():
     # Check if the file was provided in the form-data
     if 'input_image' not in request.files:
@@ -213,6 +218,7 @@ def infer():
 
 # Modify status endpoint to receive process_id in request body
 @app.route('/status', methods=['POST'])
+@spaces.GPU
 def status():
     data = request.json
     process_id = data.get('process_id')