Update app.py

app.py

@@ -2,27 +2,7 @@ import sys
 import os
 
 sys.path.append('./')
-#libraries
 os.system("pip install gradio accelerate==0.25.0 torchmetrics==1.2.1 tqdm==4.66.1 fastapi==0.111.0 transformers==4.36.2 diffusers==0.25 einops==0.7.0 bitsandbytes scipy==1.11.1 opencv-python gradio==4.24.0 fvcore cloudpickle omegaconf pycocotools basicsr av onnxruntime==1.16.2 peft==0.11.1 huggingface_hub==0.24.7 --no-deps")
-os.system("pip install ninja")
-
-#DensePose
-os.system("git clone https://github.com/facebookresearch/detectron2.git")
-os.system("!pip install git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose")
-from detectron2.config import get_cfg
-import torch
-import cv2
-import numpy as np
-from detectron2.engine import DefaultPredictor
-from densepose import add_densepose_config
-from densepose.vis.extractor import DensePoseResultExtractor
-from densepose.vis.densepose_results import DensePoseResultsFineSegmentationVisualizer as Visualizer
-# Initialize Detectron2 configuration for DensePose
-cfg = get_cfg()
-add_densepose_config(cfg)
-cfg.merge_from_file("detectron2/projects/DensePose/configs/densepose_rcnn_R_50_FPN_s1x.yaml")
-cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_s1x/165712039/model_final_162be9.pkl"
-
 import spaces
 from fastapi import FastAPI
 app = FastAPI()
@@ -40,6 +20,9 @@ from transformers import (
 )
 from diffusers import DDPMScheduler,AutoencoderKL
 from typing import List
+
+import torch
+import os
 from transformers import AutoTokenizer
 import numpy as np
 from torchvision import transforms
@@ -61,7 +44,6 @@ def pil_to_binary_mask(pil_image, threshold=0):
     return output_mask
 
 
-
 base_path = 'yisol/IDM-VTON'
 
 unet = UNet2DConditionModel.from_pretrained(
@@ -142,31 +124,11 @@ pipe = TryonPipeline.from_pretrained(
 pipe.unet_encoder = UNet_Encoder
 
 @spaces.GPU
-def start_tryon(person_img, mask_img, cloth_img, garment_des, denoise_steps):
+def start_tryon(person_img, pose_img, mask_img, cloth_img, garment_des, denoise_steps, seed):
     # Assuming device is set up (e.g., "cuda" or "cpu")
     pipe.to(device)
     pipe.unet_encoder.to(device)
 
-    cfg.MODEL.DEVICE = device
-    predictor = DefaultPredictor(cfg)
-    # Read the input image
-    height, width, _ = person_img.shape
-
-    # Process the image
-    with torch.no_grad():
-     outputs = predictor(person_img)['instances']
-
-    results = DensePoseResultExtractor()(outputs)
-    cmap = cv2.COLORMAP_VIRIDIS
-
-    # Create a base image with a black background
-    background = np.zeros((height, width, 3), dtype=np.uint8)
-
-    #Visualize the results on the black background
-    visualizer = Visualizer(alpha=1, cmap=cmap)
-    pose_img = visualizer.visualize(background, results)
-  
-
     # Resize and prepare images
     garm_img = cloth_img.convert("RGB").resize((768, 1024))
     human_img = person_img.convert("RGB").resize((768, 1024))
@@ -219,7 +181,8 @@ def start_tryon(person_img, mask_img, cloth_img, garment_des, denoise_steps):
             garm_tensor = tensor_transfrom(garm_img).unsqueeze(0).to(device, torch.float16)
 
             # Prepare the generator with optional seed
-            generator = torch.Generator(device).manual_seed(42)
+            generator = torch.Generator(device).manual_seed(seed) if seed is not None else None
+
             # Generate the virtual try-on output image
             images = pipe(
                 prompt_embeds=prompt_embeds.to(device, torch.float16),
@@ -251,6 +214,7 @@ with image_blocks as demo:
     with gr.Row():
         with gr.Column():
             person_img = gr.Image(label='Person Image', sources='upload', type="pil")
+            pose_img = gr.Image(label='Pose Image', sources='upload', type="pil")
             mask_img = gr.Image(label='Mask Image', sources='upload', type="pil")
 
         with gr.Column():
@@ -259,11 +223,12 @@ with image_blocks as demo:
 
         with gr.Column():
             denoise_steps = gr.Number(label="Denoising Steps", minimum=20, maximum=40, value=30, step=1)
+            seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42)
 
         with gr.Column():
             image_out = gr.Image(label="Output Image", elem_id="output-img", show_share_button=False)
 
     try_button = gr.Button(value="Try-on")
-    try_button.click(fn=start_tryon, inputs=[person_img, mask_img, cloth_img, garment_des, denoise_steps], outputs=[image_out], api_name='tryon')
+    try_button.click(fn=start_tryon, inputs=[person_img, pose_img, mask_img, cloth_img, garment_des, denoise_steps, seed], outputs=[image_out], api_name='tryon')
 
 image_blocks.launch()