update to dresses

app.py

@@ -1,5 +1,12 @@
-import gradio as gr
+import os
+import torch
+import numpy as np
+from typing import List
 from PIL import Image
+import gradio as gr
+
+import spaces  # Moved here to avoid CUDA initialization issues
+
 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
@@ -8,21 +15,15 @@ from transformers import (
     CLIPVisionModelWithProjection,
     CLIPTextModel,
     CLIPTextModelWithProjection,
+    AutoTokenizer  # Add this line to import AutoTokenizer
 )
-from diffusers import DDPMScheduler,AutoencoderKL
-from typing import List
-
-import torch
-import os
-from transformers import AutoTokenizer
-import spaces
-import numpy as np
+from diffusers import DDPMScheduler, AutoencoderKL
 from utils_mask import get_mask_location
 from torchvision import transforms
 import apply_net
 from preprocess.humanparsing.run_parsing import Parsing
 from preprocess.openpose.run_openpose import OpenPose
-from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation
+from detectron2.data.detection_utils import convert_PIL_to_numpy, _apply_exif_orientation
 from torchvision.transforms.functional import to_pil_image
 
 
@@ -150,7 +151,7 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
     if is_checked:
         keypoints = openpose_model(human_img.resize((384,512)))
         model_parse, _ = parsing_model(human_img.resize((384,512)))
-        mask, mask_gray = get_mask_location('hd', "upper_body", model_parse, keypoints)
+        mask, mask_gray = get_mask_location('hd', "dresses", model_parse, keypoints)
         mask = mask.resize((768,1024))
     else:
         mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))