(self)
| 72 | |
| 73 | class MagCacheTests(unittest.TestCase): |
| 74 | def setUp(self): |
| 75 | # Register standard dummy block |
| 76 | TransformerBlockRegistry.register( |
| 77 | DummyBlock, |
| 78 | TransformerBlockMetadata(return_hidden_states_index=None, return_encoder_hidden_states_index=None), |
| 79 | ) |
| 80 | # Register tuple block (Flux style) |
| 81 | TransformerBlockRegistry.register( |
| 82 | TupleOutputBlock, |
| 83 | TransformerBlockMetadata(return_hidden_states_index=0, return_encoder_hidden_states_index=1), |
| 84 | ) |
| 85 | |
| 86 | def _set_context(self, model, context_name): |
| 87 | """Helper to set context on all hooks in the model.""" |
nothing calls this directly
no test coverage detected