| import torch |
|
|
| from opensora.registry import MODELS |
|
|
|
|
| @MODELS.register_module("classes") |
| class ClassEncoder: |
| def __init__(self, num_classes, model_max_length=None, device="cuda", dtype=torch.float): |
| self.num_classes = num_classes |
| self.y_embedder = None |
|
|
| self.model_max_length = model_max_length |
| self.output_dim = None |
| self.device = device |
|
|
| def encode(self, text): |
| return dict(y=torch.tensor([int(t) for t in text]).to(self.device)) |
|
|
| def null(self, n): |
| return torch.tensor([self.num_classes] * n).to(self.device) |
|
|
