Spaces:
Running on Zero
Running on Zero
| from typing import Callable, List, Sequence, Dict, Type | |
| from torch import nn, Tensor | |
| import torch | |
| from conditioning.condition_type import ConditionType | |
| from conditioning.conditioning_method import ConditioningMethod | |
| from conditioning.embedded_condition import EmbeddedCondition | |
| class ConditionDispatcher(nn.Module): | |
| def __init__(self, cond_to_method: Dict[ConditionType, ConditioningMethod], | |
| embedding_dim: int, condition_dropout: float): | |
| super().__init__() | |
| self.embedding_dim: int = embedding_dim | |
| self.condition_dropout: float = condition_dropout | |
| # what fusing method to use for each conditioning type | |
| self.cond_to_method: Dict[ConditionType, | |
| ConditioningMethod] = cond_to_method | |
| # what condition types is each method applied to | |
| self.method_to_conds: Dict[ConditioningMethod, List[ConditionType]] = {} | |
| for c, m in self.cond_to_method.items(): | |
| self.method_to_conds[m] = self.method_to_conds.get(m, []) + [c] | |
| # instantiate fuser classes | |
| self.method_to_fuser: nn.ModuleDict = nn.ModuleDict({ | |
| m.value: | |
| METHOD_TO_FUSER_TYPE[m](len(self.method_to_conds[m]), | |
| self.embedding_dim) | |
| for m in self.method_to_conds.keys() | |
| }) | |
| def dropout(self, x: EmbeddedCondition): | |
| if self.training and self.condition_dropout > 0: | |
| drop = torch.rand(x.data.shape[0]) < self.condition_dropout | |
| # x.data[drop] = torch.zeros_like(x.data[drop]) #.detach() | |
| x.data[drop] = x.data[drop] * 0 | |
| if x.mask is not None: | |
| x.mask[drop] = torch.zeros_like(x.mask[drop]) #.detach() | |
| return x | |
| def forward( | |
| self, embedded_conditions: Dict[ConditionType, EmbeddedCondition] | |
| ) -> Dict[ConditioningMethod, EmbeddedCondition]: | |
| assert embedded_conditions.keys() == self.cond_to_method.keys() | |
| # figure out where does each conditioning go, between: | |
| # - summed to the input of the lm (after embedding) | |
| # - prepended to the input of the lm (after embedding) | |
| # - in cross-attention of the lm | |
| method_to_embeddings: Dict[ConditioningMethod, | |
| List[EmbeddedCondition]] = {} | |
| # for condition_type, emb in embedded_conditions.items(): | |
| for condition_type in sorted(embedded_conditions.keys()): | |
| emb = embedded_conditions[condition_type] | |
| # apply dropout to each condition | |
| emb = self.dropout(emb) | |
| method: ConditioningMethod = self.cond_to_method[condition_type] | |
| method_to_embeddings[method] = method_to_embeddings.get(method, | |
| []) + [emb] | |
| # has to return a dict that associates one embedded condition to each | |
| # conditioning method | |
| method_to_fused_embedding: Dict[ConditioningMethod, | |
| EmbeddedCondition] = {} | |
| # if multiple conditions are dispatched to the same method, we have to | |
| # fuse them | |
| for m, emb_list in method_to_embeddings.items(): | |
| fused_embedding: EmbeddedCondition = self.method_to_fuser[m.value]( | |
| emb_list) | |
| method_to_fused_embedding[m] = fused_embedding | |
| return method_to_fused_embedding | |
| class CrossAttentionFuser(nn.Module): | |
| def __init__(self, n_conditions: int, embedding_dim: int): | |
| super().__init__() | |
| self.n_conditions: int = n_conditions | |
| self.embedding_dim: int = embedding_dim | |
| # use concatenation with segment embedding to merge conditions | |
| if self.n_conditions > 1: | |
| self.segment_embedding = nn.Embedding(self.n_conditions - 1, | |
| self.embedding_dim) | |
| def forward(self, conds: Sequence[EmbeddedCondition]) -> EmbeddedCondition: | |
| if len(conds) == 0: | |
| raise RuntimeError("Received a list of 0 length. There is a bug.") | |
| if len(conds) == 1: | |
| return conds[0] | |
| # apply segment embedding | |
| if len(conds) == 2: | |
| assert self.n_conditions == 2 | |
| assert conds[0].mask is not None and conds[1].mask is not None | |
| cond1_emb = conds[1].data | |
| cond1_ids = torch.zeros(cond1_emb.shape[0], | |
| cond1_emb.shape[1], | |
| dtype=torch.long, | |
| device=conds[1].data.device) | |
| cond1_segembed = self.segment_embedding(cond1_ids) | |
| cond1_emb = cond1_emb + cond1_segembed | |
| total_cond = torch.cat((conds[0].data, cond1_emb), dim=1) | |
| total_mask = torch.cat((conds[0].mask, conds[1].mask), dim=1) | |
| return EmbeddedCondition(total_cond, total_mask) | |
| # multiple conditions in cross-atteniton are not implemented | |
| raise NotImplementedError() | |
| class PrependFuser(nn.Module): | |
| def __init__(self, n_conditions: int, embedding_dim: int): | |
| super().__init__() | |
| self.n_conditions: int = n_conditions | |
| self.embedding_dim: int = embedding_dim | |
| def forward(self, conds: Sequence[EmbeddedCondition]): | |
| if len(conds) == 0: | |
| raise RuntimeError("Received a list of 0 length. There is a bug.") | |
| if len(conds) == 1: | |
| return conds[0] | |
| # multiple conditions prepended to input are not implemented | |
| raise NotImplementedError() | |
| class SumFuser(nn.Module): | |
| def __init__(self, n_conditions: int, embedding_dim: int): | |
| super().__init__() | |
| self.n_conditions: int = n_conditions | |
| self.embedding_dim: int = embedding_dim | |
| def forward(self, conds: Sequence[EmbeddedCondition]): | |
| if len(conds) == 0: | |
| raise RuntimeError("Received a list of 0 length. There is a bug.") | |
| if len(conds) == 1: | |
| return conds[0] | |
| # multiple conditions summed to input are not implemented | |
| raise NotImplementedError() | |
| METHOD_TO_FUSER_TYPE: Dict[ConditioningMethod, Type[nn.Module]] = { | |
| ConditioningMethod.CROSS_ATTENTION: CrossAttentionFuser, | |
| ConditioningMethod.INPUT_PREPEND: PrependFuser, | |
| ConditioningMethod.INPUT_SUM: SumFuser, | |
| } | |