add v2 infos

README.md

@@ -15,48 +15,44 @@ pipeline_tag: image-to-image
 library_name: diffusers
 ---
 
-## TryOffDiff
+## TryOffDiff <small>[[Project Page]][project_page]</small>
 
-The models proposed in the paper _"TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models"_
-[[paper]][paper_arxiv] [[project page]][project_page]:
-- `tryoffdiff.pth`: The pre-trained StableDiffusion-v1.4 fine-tuned on `VITON-HD-train` dataset.
+These models are based on the research papers:
+1. _"TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models"_ [[paper]][paper_arxiv]
+2. _"Enhancing Person-to-Person Virtual Try-On with Multi-Garment Virtual Try-Off"_ [[paper]][paper2_arxiv]
 
+Available models:
+- `tryoffdiff.pth`: Generates upper-body garments. Fine-tuned on `VITON-HD-train`.
+- `tryoffdiffv2_upper.pth`: Generates upper-body garments. Fine-tuned on a subset of `Dress Code-train`.
+- `tryoffdiffv2_lower.pth`: Generates lower-body garments. Fine-tuned on a subset of `Dress Code-train`.
+- `tryoffdiffv2_dress.pth`: Generates dresses. Fine-tuned on a subset of `Dress Code-train`.
+- `tryoffdiffv2_multi.pth`: Generates all three garment types (upper-body, lower-body, dresses). Fine-tuned on the entire `Dress Code-train`.
 
 ## Usage
 
 ```python
+import torch
 from huggingface_hub import hf_hub_download
+from tryoffdiff.modeling.model import create_model
 
-class TryOffDiff(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
-        self.transformer = torch.nn.TransformerEncoderLayer(d_model=768, nhead=8, batch_first=True)
-        self.proj = nn.Linear(1024, 77)
-        self.norm = nn.LayerNorm(768)
+repo = "rizavelioglu/tryoffdiff"  # Define the repository
 
-    def adapt_embeddings(self, x):
-        x = self.transformer(x)
-        x = self.proj(x.permute(0, 2, 1)).permute(0, 2, 1)
-        return self.norm(x)
+# --- Helper function to load a model ---
+def load_tryoffdiff_model(model_name, filename, device="cuda"):
+    model = create_model(model_name)
+    model_path = hf_hub_download(repo_id=repo, filename=filename)
+    model.load_state_dict(torch.load(model_path, weights_only=True))
+    return model.eval().to(device)
 
-    def forward(self, noisy_latents, t, cond_emb):
-        cond_emb = self.adapt_embeddings(cond_emb)
-        return self.unet(noisy_latents, t, encoder_hidden_states=cond_emb).sample
-
-path_model = hf_hub_download(
-    repo_id="rizavelioglu/tryoffdiff",
-    filename="tryoffdiff.pth",  # or one of ablations ["ldm-1", "ldm-2", "ldm-3", ...]
-)
-net = TryOffDiff()
-net.load_state_dict(torch.load(path_model, weights_only=False))
-net.eval().to(device)
+tryoffdiffv1       = load_tryoffdiff_model("TryOffDiff", "tryoffdiff.pth")
+tryoffdiffv2_upper = load_tryoffdiff_model("TryOffDiffv2_single", "tryoffdiffv2_upper.pth")
+tryoffdiffv2_lower = load_tryoffdiff_model("TryOffDiffv2_single", "tryoffdiffv2_lower.pth")
+tryoffdiffv2_dress = load_tryoffdiff_model("TryOffDiffv2_single", "tryoffdiffv2_dress.pth")
+tryoffdiffv2_multi = load_tryoffdiff_model("TryOffDiffv2", "tryoffdiffv2_multi.pth")
 ```
 
-> Check out the demo code on [HuggingFace Spaces][hf_spaces] for the full running example.
-
-> Also, check out [GitHub repository][github] to get more information on
-> training, inference, and evaluation.
+> For a complete running example, see the demo code on [HuggingFace Spaces][hf_spaces].\
+> For more information on training, inference, and evaluation, refer to the [GitHub repository][github].
 
 ### License
 TL;DR: Not available for commercial use, unless the FULL source code is shared! \
@@ -64,7 +60,7 @@ This project is intended solely for academic research. No commercial benefits ar
 Models are licensed under [Server Side Public License (SSPL)][license]
 
 ### Citation
-If you find this repository useful in your research, please consider giving a star ⭐ and a citation:
+If you use this work, please give a star ⭐ and a citation:
 ```
 @article{velioglu2024tryoffdiff,
   title     = {TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models},
@@ -73,11 +69,19 @@ If you find this repository useful in your research, please consider giving a st
   year      = {2024},
   note      = {\url{https://doi.org/nt3n}}
 }
+@article{velioglu2025enhancing,
+  title     = {Enhancing Person-to-Person Virtual Try-On with Multi-Garment Virtual Try-Off},
+  author    = {Velioglu, Riza and Bevandic, Petra and Chan, Robin and Hammer, Barbara},
+  journal   = {arXiv},
+  year      = {2025},
+  note      = {\url{https://doi.org/pn67}}
+}
 ```
 
 [hf_spaces]: https://huggingface.co/spaces/rizavelioglu/tryoffdiff/blob/main/app.py
 [project_page]: https://rizavelioglu.github.io/tryoffdiff/
 [paper_arxiv]: https://arxiv.org/abs/2411.18350
+[paper2_arxiv]: https://arxiv.org/abs/2504.13078
 [github]: https://github.com/rizavelioglu/tryoffdiff
 [license]: https://www.mongodb.com/legal/licensing/server-side-public-license