(self, cond, force_mask=False)
| 97 | return clip_model |
| 98 | |
| 99 | def mask_cond(self, cond, force_mask=False): |
| 100 | bs = cond.shape[0] |
| 101 | if force_mask: |
| 102 | return torch.zeros_like(cond) |
| 103 | elif self.training and self.cond_mask_prob > 0.: |
| 104 | mask = torch.ones(bs, device=cond.device) * self.cond_mask_prob |
| 105 | # 1-> use null_cond, 0-> use real cond |
| 106 | mask = torch.bernoulli(mask).view(bs, 1) |
| 107 | return cond * (1. - mask) |
| 108 | else: |
| 109 | return cond |
| 110 | |
| 111 | def encode_text(self, raw_text): |
| 112 | device = next(self.parameters()).device |