Update app.py

app.py

@@ -7,8 +7,9 @@ import torch
 import torchvision.transforms as transforms
 from transformers import pipeline
 
-# Load the DeepVTO model from Hugging Face
-deepvto_pipeline = pipeline("image-to-image", model="huggingface/deepvto")
+# Load the Rao's Virtual Try-On model from Hugging Face
+model_id = "gouthaml/raos-virtual-try-on-model"
+tryon_pipeline = pipeline("image-to-image", model=model_id)
 
 # Load sample product data
 product_data = pd.DataFrame({
@@ -31,16 +32,12 @@ def process_image(image, product):
     
     garment_image = Image.open(garment_path).convert("RGB")
     
-    # Convert images to the format required by the model
-    person_image_tensor = transforms.ToTensor()(person_image).unsqueeze(0)
-    garment_image_tensor = transforms.ToTensor()(garment_image).unsqueeze(0)
+    # Run the virtual try-on model
+    inputs = {"image": person_image, "garment": garment_image}
+    result = tryon_pipeline(inputs)
     
-    # Run the DeepVTO model
-    with torch.no_grad():
-        output = deepvto_pipeline(person_image_tensor, garment_image_tensor)
-    
-    # Convert the output to a PIL image
-    output_image = transforms.ToPILImage()(output[0])
+    # Post-process the output
+    output_image = result[0]  # Assuming the model returns a list of outputs
     
     # Convert to numpy array for Gradio
     result_array = np.array(output_image)