Update app.py

app.py

@@ -7,6 +7,7 @@ import base64
 import uuid
 import random
 import logging
+import spaces
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
 from src.unet_hacked_tryon import UNet2DConditionModel
@@ -154,6 +155,7 @@ def clear_gpu_memory():
 
 # Main try-on function
 @torch.no_grad()
+@spaces.GPU
 def start_tryon(human_dict, garment_img, garment_des, use_auto_mask, use_auto_crop, denoise_steps, seed, categorie='upper_body'):
     try:
         device = torch.device("cuda")
@@ -208,7 +210,7 @@ def start_tryon(human_dict, garment_img, garment_des, use_auto_mask, use_auto_cr
         logging.error(f"Error during try-on: {e}")
         raise
     finally:
-        clear_gpu_memory()
+        #clear_gpu_memory()
 
 # API endpoints
 @app.route('/tryon', methods=['POST'])