Update app.py

app.py

@@ -198,8 +198,13 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
         with torch.cuda.amp.autocast():
             with torch.no_grad():
                 # 5.1 生成正面提示词嵌入
-                prompt = "((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), the model is wearing " + garment_des
-                negative_prompt = "monochrome, lowres, bad anatomy, worst quality, normal quality, low quality, blurry, jpeg artifacts, sketch"
+                if category == "saree":
+                    prompt = "((best quality, ultra detailed, masterpiece)), a realistic Indian woman wearing a traditional saree with " + garment_des
+                    negative_prompt = "monochrome, lowres, bad anatomy, worst quality, normal quality, low quality, blurry, jpeg artifacts, sketch"
+                else:
+                    prompt = "((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), the model is wearing " + garment_des
+                    negative_prompt = "monochrome, lowres, bad anatomy, worst quality, normal quality, low quality, blurry, jpeg artifacts, sketch"
+                
                 with torch.inference_mode():
                     # 编码提示词
                     (
@@ -307,7 +312,7 @@ with image_blocks as demo:
                 is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
             with gr.Row():
                 category = gr.Dropdown(
-                    choices=["upper_body", "lower_body", "dresses"],
+                    choices=["upper_body", "lower_body", "dresses", "saree"],
                     label="Category",
                     value="upper_body"
                 )