cuda

app.py

@@ -14,6 +14,8 @@ import torch
 from omegaconf import OmegaConf
 from PIL import Image
 import spaces
+print(torch.cuda.is_available(), torch.cuda.device_count())
+
 
 from cldm.model import create_model
 from cldm.plms_hacked import PLMSSampler
@@ -172,6 +174,7 @@ def process_hd(vton_img, garm_img, n_steps, is_custom):
     vton_img = vton_img.resize((IMG_W, IMG_H))  # size for densepose
     densepose = densepose_model_hd.execute(vton_img)  # densepose
     print('%.2fs' % (time.time() - stt))
+    breakpoint()
 
     batch = get_batch(
         vton_img, 

utils_stableviton.py

@@ -176,7 +176,7 @@ def get_tensor(img, h, w, is_mask=False):
         img = (img / 127.5) - 1.0
     else:
         img = (img < 128).astype(np.float32)[:,:,None]
-    return torch.from_numpy(img)[None]
+    return torch.from_numpy(img)[None].cuda()
     
 def get_batch(image, cloth, densepose, agn_img, agn_mask, img_h, img_w):
     batch = dict()