Update app.py

app.py

@@ -9,7 +9,15 @@ from cloth_segmentation.networks.u2net import U2NET  # Import U²-Net
 # Load U²-Net model
 model_path = "cloth_segmentation/networks/cloth_segm_u2net_latest.pth"  # Ensure this path is correct
 model = U2NET(3, 1)
-model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+
+# Load the state dictionary
+state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+
+# Remove the 'module.' prefix from the keys
+state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+
+# Load the modified state dictionary into the model
+model.load_state_dict(state_dict)
 model.eval()
 
 def segment_dress(image_np):