Update README.md

README.md

@@ -29,14 +29,33 @@ The models proposed in the paper _"TryOffDiff: Virtual-Try-Off via High-Fidelity
 ```python
 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 ["ldm-1", "ldm-2", "ldm-3", ...]
+    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][hf_spaces] for visualizing the output.
+> 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.
@@ -58,7 +77,7 @@ If you find this repository useful in your research, please consider giving a st
 }
 ```
 
-[hf_spaces]: https://huggingface.co/spaces/rizavelioglu/tryoffdiff
+[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
 [github]: https://github.com/rizavelioglu/tryoffdiff