| 96 | |
| 97 | |
| 98 | class LTXModifiedBasicTransformerBlock(BasicTransformerBlock): |
| 99 | def forward( |
| 100 | self, |
| 101 | x, |
| 102 | context=None, |
| 103 | attention_mask=None, |
| 104 | timestep=None, |
| 105 | pe=None, |
| 106 | transformer_options={}, |
| 107 | ): |
| 108 | shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( |
| 109 | self.scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) |
| 110 | + timestep.reshape( |
| 111 | x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1 |
| 112 | ) |
| 113 | ).unbind(dim=2) |
| 114 | x += ( |
| 115 | self.attn1( |
| 116 | comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa, |
| 117 | pe=pe, |
| 118 | transformer_options=transformer_options, |
| 119 | ) |
| 120 | * gate_msa |
| 121 | ) |
| 122 | |
| 123 | x += self.attn2(x, context=context, mask=attention_mask) |
| 124 | |
| 125 | y = comfy.ldm.common_dit.rms_norm(x) * (1 + scale_mlp) + shift_mlp |
| 126 | x += self.ff(y) * gate_mlp |
| 127 | |
| 128 | return x |
| 129 | |
| 130 | |
| 131 | class LTXVModelModified(LTXVModel): |
nothing calls this directly
no outgoing calls
no test coverage detected