Drops labels to enable classifier-free guidance.
(self, caption, force_drop_ids=None)
| 202 | self.uncond_prob = uncond_prob |
| 203 | |
| 204 | def token_drop(self, caption, force_drop_ids=None): |
| 205 | """ |
| 206 | Drops labels to enable classifier-free guidance. |
| 207 | """ |
| 208 | if force_drop_ids is None: |
| 209 | drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob |
| 210 | else: |
| 211 | drop_ids = force_drop_ids == 1 |
| 212 | caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, caption) |
| 213 | return caption |
| 214 | |
| 215 | def forward(self, caption, train, force_drop_ids=None): |
| 216 | if train: |