Update app.py

app.py

@@ -5,13 +5,12 @@ from os import path
 from safetensors.torch import load_file
 import huggingface_hub
 from huggingface_hub import hf_hub_download
-
-
+import os
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
 os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
-
+import spaces
 import gradio as gr
 import torch
 from diffusers import FluxPipeline
@@ -92,7 +91,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=4):
             output = gr.Image(label="Your Generated Image")
             
-
+    @spaces.GPU
     def process_image(height, width, steps, scales, prompt, seed, preset_prompt):
         global pipe
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):