Spaces:
Runtime error
Runtime error
| # Copyright (c) 2023-2024, Zexin He | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # https://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| from functools import partial | |
| import torch | |
| import torch.nn as nn | |
| from accelerate.logging import get_logger | |
| from typing import Any, Dict, Optional, Tuple, Union | |
| from diffusers.utils import is_torch_version | |
| logger = get_logger(__name__) | |
| class TransformerDecoder(nn.Module): | |
| """ | |
| Transformer blocks that process the input and optionally use condition and modulation. | |
| """ | |
| def __init__(self, block_type: str, | |
| num_layers: int, num_heads: int, | |
| inner_dim: int, cond_dim: int = None, mod_dim: int = None, | |
| gradient_checkpointing=False, | |
| eps: float = 1e-6, | |
| use_dual_attention: bool = False,): | |
| super().__init__() | |
| self.gradient_checkpointing = gradient_checkpointing | |
| self.block_type = block_type | |
| if block_type == "sd3_cond": | |
| # dual_attention_layers = list(range(num_layers//2)) | |
| dual_attention_layers = [] | |
| self.layers = nn.ModuleList([ | |
| self._block_fn(inner_dim, cond_dim, mod_dim)( | |
| num_heads=num_heads, | |
| eps=eps, | |
| context_pre_only=i == num_layers - 1, | |
| use_dual_attention=use_dual_attention, # True if i in dual_attention_layers else False, | |
| ) | |
| for i in range(num_layers) | |
| ]) | |
| else: | |
| self.layers = nn.ModuleList([ | |
| self._block_fn(inner_dim, cond_dim, mod_dim)( | |
| num_heads=num_heads, | |
| eps=eps, | |
| ) | |
| for _ in range(num_layers) | |
| ]) | |
| self.norm = nn.LayerNorm(inner_dim, eps=eps) | |
| if self.block_type in ["cogvideo_cond", "sd3_cond"]: | |
| self.linear_cond_proj = nn.Linear(cond_dim, inner_dim) | |
| def block_type(self): | |
| return self._block_type | |
| def block_type(self, block_type): | |
| assert block_type in ['basic', 'cond', 'mod', 'cond_mod', 'sd3_cond', 'cogvideo_cond'], \ | |
| f"Unsupported block type: {block_type}" | |
| self._block_type = block_type | |
| def _block_fn(self, inner_dim, cond_dim, mod_dim): | |
| assert inner_dim is not None, f"inner_dim must always be specified" | |
| if self.block_type == 'basic': | |
| assert cond_dim is None and mod_dim is None, \ | |
| f"Condition and modulation are not supported for BasicBlock" | |
| from .block import BasicBlock | |
| # logger.debug(f"Using BasicBlock") | |
| return partial(BasicBlock, inner_dim=inner_dim) | |
| elif self.block_type == 'cond': | |
| assert cond_dim is not None, f"Condition dimension must be specified for ConditionBlock" | |
| assert mod_dim is None, f"Modulation dimension is not supported for ConditionBlock" | |
| from .block import ConditionBlock | |
| # logger.debug(f"Using ConditionBlock") | |
| return partial(ConditionBlock, inner_dim=inner_dim, cond_dim=cond_dim) | |
| elif self.block_type == 'mod': | |
| # logger.error(f"modulation without condition is not implemented") | |
| raise NotImplementedError(f"modulation without condition is not implemented") | |
| elif self.block_type == 'cond_mod': | |
| assert cond_dim is not None and mod_dim is not None, \ | |
| f"Condition and modulation dimensions must be specified for ConditionModulationBlock" | |
| from .block import ConditionModulationBlock | |
| # logger.debug(f"Using ConditionModulationBlock") | |
| return partial(ConditionModulationBlock, inner_dim=inner_dim, cond_dim=cond_dim, mod_dim=mod_dim) | |
| elif self.block_type == 'cogvideo_cond': | |
| # logger.debug(f"Using CogVideoXBlock") | |
| from lam.models.transformer_dit import CogVideoXBlock | |
| # assert inner_dim == cond_dim, f"inner_dim:{inner_dim}, cond_dim:{cond_dim}" | |
| return partial(CogVideoXBlock, dim=inner_dim, attention_bias=True) | |
| elif self.block_type == 'sd3_cond': | |
| # logger.debug(f"Using SD3JointTransformerBlock") | |
| from lam.models.transformer_dit import SD3JointTransformerBlock | |
| return partial(SD3JointTransformerBlock, dim=inner_dim, qk_norm="rms_norm") | |
| else: | |
| raise ValueError(f"Unsupported block type during runtime: {self.block_type}") | |
| def assert_runtime_integrity(self, x: torch.Tensor, cond: torch.Tensor, mod: torch.Tensor): | |
| assert x is not None, f"Input tensor must be specified" | |
| if self.block_type == 'basic': | |
| assert cond is None and mod is None, \ | |
| f"Condition and modulation are not supported for BasicBlock" | |
| elif 'cond' in self.block_type: | |
| assert cond is not None and mod is None, \ | |
| f"Condition must be specified and modulation is not supported for ConditionBlock" | |
| elif self.block_type == 'mod': | |
| raise NotImplementedError(f"modulation without condition is not implemented") | |
| else: | |
| assert cond is not None and mod is not None, \ | |
| f"Condition and modulation must be specified for ConditionModulationBlock" | |
| def forward_layer(self, layer: nn.Module, x: torch.Tensor, cond: torch.Tensor, mod: torch.Tensor): | |
| if self.block_type == 'basic': | |
| return layer(x) | |
| elif self.block_type == 'cond': | |
| return layer(x, cond) | |
| elif self.block_type == 'mod': | |
| return layer(x, mod) | |
| else: | |
| return layer(x, cond, mod) | |
| def forward(self, x: torch.Tensor, cond: torch.Tensor = None, mod: torch.Tensor = None): | |
| # x: [N, L, D] | |
| # cond: [N, L_cond, D_cond] or None | |
| # mod: [N, D_mod] or None | |
| self.assert_runtime_integrity(x, cond, mod) | |
| if self.block_type in ["cogvideo_cond", "sd3_cond"]: | |
| cond = self.linear_cond_proj(cond) | |
| for layer in self.layers: | |
| if self.training and self.gradient_checkpointing: | |
| def create_custom_forward(module): | |
| def custom_forward(*inputs): | |
| return module(*inputs) | |
| return custom_forward | |
| ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} | |
| x, cond = torch.utils.checkpoint.checkpoint( | |
| create_custom_forward(layer), | |
| x, | |
| cond, | |
| **ckpt_kwargs, | |
| ) | |
| else: | |
| x, cond = layer( | |
| hidden_states=x, | |
| encoder_hidden_states=cond, | |
| temb=None, | |
| # image_rotary_emb=None, | |
| ) | |
| x = self.norm(x) | |
| else: | |
| for layer in self.layers: | |
| x = self.forward_layer(layer, x, cond, mod) | |
| x = self.norm(x) | |
| return x | |