(
self,
resolution: int,
in_channels: int,
model_channels: int,
cond_channels: int,
out_channels: int,
num_blocks: int,
num_heads: Optional[int] = None,
num_head_channels: Optional[int] = 64,
mlp_ratio: float = 4,
patch_size: int = 2,
num_io_res_blocks: int = 2,
io_block_channels: List[int] = None,
pe_mode: Literal["ape", "rope"] = "ape",
use_fp16: bool = False,
use_checkpoint: bool = False,
use_skip_connection: bool = True,
share_mod: bool = False,
qk_rms_norm: bool = False,
qk_rms_norm_cross: bool = False,
)
| 68 | |
| 69 | class SLatFlowModel(nn.Module): |
| 70 | def __init__( |
| 71 | self, |
| 72 | resolution: int, |
| 73 | in_channels: int, |
| 74 | model_channels: int, |
| 75 | cond_channels: int, |
| 76 | out_channels: int, |
| 77 | num_blocks: int, |
| 78 | num_heads: Optional[int] = None, |
| 79 | num_head_channels: Optional[int] = 64, |
| 80 | mlp_ratio: float = 4, |
| 81 | patch_size: int = 2, |
| 82 | num_io_res_blocks: int = 2, |
| 83 | io_block_channels: List[int] = None, |
| 84 | pe_mode: Literal["ape", "rope"] = "ape", |
| 85 | use_fp16: bool = False, |
| 86 | use_checkpoint: bool = False, |
| 87 | use_skip_connection: bool = True, |
| 88 | share_mod: bool = False, |
| 89 | qk_rms_norm: bool = False, |
| 90 | qk_rms_norm_cross: bool = False, |
| 91 | ): |
| 92 | super().__init__() |
| 93 | self.resolution = resolution |
| 94 | self.in_channels = in_channels |
| 95 | self.model_channels = model_channels |
| 96 | self.cond_channels = cond_channels |
| 97 | self.out_channels = out_channels |
| 98 | self.num_blocks = num_blocks |
| 99 | self.num_heads = num_heads or model_channels // num_head_channels |
| 100 | self.mlp_ratio = mlp_ratio |
| 101 | self.patch_size = patch_size |
| 102 | self.num_io_res_blocks = num_io_res_blocks |
| 103 | self.io_block_channels = io_block_channels |
| 104 | self.pe_mode = pe_mode |
| 105 | self.use_fp16 = use_fp16 |
| 106 | self.use_checkpoint = use_checkpoint |
| 107 | self.use_skip_connection = use_skip_connection |
| 108 | self.share_mod = share_mod |
| 109 | self.qk_rms_norm = qk_rms_norm |
| 110 | self.qk_rms_norm_cross = qk_rms_norm_cross |
| 111 | self.dtype = torch.float16 if use_fp16 else torch.float32 |
| 112 | |
| 113 | if self.io_block_channels is not None: |
| 114 | assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2" |
| 115 | assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages" |
| 116 | |
| 117 | self.t_embedder = TimestepEmbedder(model_channels) |
| 118 | if share_mod: |
| 119 | self.adaLN_modulation = nn.Sequential( |
| 120 | nn.SiLU(), |
| 121 | nn.Linear(model_channels, 6 * model_channels, bias=True) |
| 122 | ) |
| 123 | |
| 124 | if pe_mode == "ape": |
| 125 | self.pos_embedder = AbsolutePositionEmbedder(model_channels) |
| 126 | |
| 127 | self.input_layer = sp.SparseLinear(in_channels, model_channels if io_block_channels is None else io_block_channels[0]) |
no test coverage detected