Grok Fix add @spaces

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import numpy as np
 import torch
@@ -5,14 +6,6 @@ from diffusers import FluxFillPipeline
 from PIL import Image
 import random
 
-# Set device (GPU if available, otherwise CPU)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load the pipeline
-pipe = FluxFillPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
-).to(device)
-
 
 # Function to calculate optimal dimensions for the output image
 def calculate_optimal_dimensions(image: Image.Image):
@@ -43,10 +36,16 @@ def calculate_optimal_dimensions(image: Image.Image):
 
 
 # Inpainting function
+@spaces.GPU
 def infer(image: Image.Image, prompt: str):
     if image is None:
         raise gr.Error("Please upload an image.")
 
+    # Load the pipeline inside the function on "cuda"
+    pipe = FluxFillPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
+    ).to("cuda")
+
     # Calculate optimal dimensions
     width, height = calculate_optimal_dimensions(image)
 
@@ -56,7 +55,7 @@ def infer(image: Image.Image, prompt: str):
 
     # Set seed for reproducibility
     seed = random.randint(0, np.iinfo(np.int32).max)
-    generator = torch.Generator(device).manual_seed(seed)
+    generator = torch.Generator("cuda").manual_seed(seed)
 
     try:
         output = pipe(