rizavelioglu/tryoffdiff

TryOffDiff

The models proposed in the paper "TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models" [paper] [project page]:

• tryoffdiff.pth: The pre-trained StableDiffusion-v1.4 fine-tuned on VITON-HD-train dataset.
• .pth: A U-Net trained from scratch on VITON-HD-train dataset.
• .pth:

Usage

from huggingface_hub import hf_hub_download

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)

    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)

    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)

Check out the demo code on HuggingFace Spaces for the full running example.

Also, check out GitHub repository to get more information on training, inference, and evaluation.

License

TL;DR: Not available for commercial use, unless the FULL source code is shared!
This project is intended solely for academic research. No commercial benefits are derived from it. Models are licensed under Server Side Public License (SSPL)

Citation

If you find this repository useful in your research, please consider giving a star ⭐ and a citation:

@article{velioglu2024tryoffdiff,
  title     = {TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models},
  author    = {Velioglu, Riza and Bevandic, Petra and Chan, Robin and Hammer, Barbara},
  journal   = {arXiv},
  year      = {2024},
  note      = {\url{https://doi.org/nt3n}}
}