Update app.py

app.py

@@ -1,11 +1,15 @@
 import argparse
 import os
+os.environ['CUDA_HOME'] = '/usr/local/cuda'
+os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
 from datetime import datetime
 
+
+import spaces
+
 import gradio as gr
 import numpy as np
 import torch
-device = torch.device('cpu')  # Explicitly use CPU if desired
 
 from diffusers.image_processor import VaeImageProcessor
 from huggingface_hub import snapshot_download
@@ -111,18 +115,19 @@ pipeline = CatVTONPipeline(
     attn_ckpt_version="mix",
     weight_dtype=init_weight_dtype(args.mixed_precision),
     use_tf32=args.allow_tf32,
-    # device='cuda'
-    device='cpu'
+    device='cuda'
+    #device='cpu'
 )
 # AutoMasker
 mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True)
 automasker = AutoMasker(
     densepose_ckpt=os.path.join(repo_path, "DensePose"),
     schp_ckpt=os.path.join(repo_path, "SCHP"),
-    # device='cuda', 
-    device='cpu'
+    device='cuda', 
+    #device='cpu'
 )
 
+@spaces.GPU(duration=120)
 def submit_function(
     person_image,
     cloth_image,
@@ -149,8 +154,8 @@ def submit_function(
 
     generator = None
     if seed != -1:
-        # generator = torch.Generator(device='cuda').manual_seed(seed)
-        generator = torch.Generator(device='cpu').manual_seed(seed)
+        generator = torch.Generator(device='cuda').manual_seed(seed)
+        #generator = torch.Generator(device='cpu').manual_seed(seed)
 
     person_image = Image.open(person_image).convert("RGB")
     cloth_image = Image.open(cloth_image).convert("RGB")