Made functions use gpu

app.py

@@ -68,16 +68,18 @@ def init_model():
         dtype=torch.float32
     )
     
-    return model, pipeline
+    return model, pipeline, args
 
 # Global variables to store model and pipeline
 model = None
 pipeline = None
 
+@spaces.GPU
 def get_model():
-    global model, pipeline
+    """Lazy initialization of model and pipeline"""
+    global model, pipeline, args
     if model is None or pipeline is None:
-        model, pipeline = init_model()
+        model, pipeline, args = init_model()
     return model, pipeline
 
 rembg_session = rembg.new_session()
@@ -144,6 +146,7 @@ def add_random_background(image, color):
     background = Image.new("RGBA", image.size, color)
     return Image.alpha_composite(background, image)
 
+@spaces.GPU
 def preprocess_image(input_image, background_choice, foreground_ratio, back_groud_color):
     """Preprocess the input image"""
     try:
@@ -169,6 +172,7 @@ def preprocess_image(input_image, background_choice, foreground_ratio, back_grou
         print(f"Error in preprocess_image: {str(e)}")
         raise e
 
+@spaces.GPU
 def gen_image(processed_image, seed, scale, step):
     """Generate the 3D model"""
     try: