Update ootd/inference_ootd.py

ootd/inference_ootd.py

@@ -41,7 +41,12 @@ import spaces
 @spaces.GPU
 def process_hd(vton_img, garm_img, n_samples, n_steps, image_scale, seed):
     model_type = 'hd'
-    category = 0 # 0:upperbody; 1:lowerbody; 2:dress
+     if category == 'Upper-body':
+        category = 0
+    elif category == 'Lower-body':
+        category = 1
+    else:
+        category =2
 
     with torch.no_grad():
         openpose_model_hd.preprocessor.body_estimation.model.to('cuda')