| from torch.utils.checkpoint import checkpoint |
|
|
| import ldm.modules.attention |
| import ldm.modules.diffusionmodules.openaimodel |
|
|
|
|
| def BasicTransformerBlock_forward(self, x, context=None): |
| return checkpoint(self._forward, x, context) |
|
|
|
|
| def AttentionBlock_forward(self, x): |
| return checkpoint(self._forward, x) |
|
|
|
|
| def ResBlock_forward(self, x, emb): |
| return checkpoint(self._forward, x, emb) |
|
|
|
|
| stored = [] |
|
|
|
|
| def add(): |
| if len(stored) != 0: |
| return |
|
|
| stored.extend([ |
| ldm.modules.attention.BasicTransformerBlock.forward, |
| ldm.modules.diffusionmodules.openaimodel.ResBlock.forward, |
| ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward |
| ]) |
|
|
| ldm.modules.attention.BasicTransformerBlock.forward = BasicTransformerBlock_forward |
| ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = ResBlock_forward |
| ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = AttentionBlock_forward |
|
|
|
|
| def remove(): |
| if len(stored) == 0: |
| return |
|
|
| ldm.modules.attention.BasicTransformerBlock.forward = stored[0] |
| ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = stored[1] |
| ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = stored[2] |
|
|
| stored.clear() |
|
|
|
|
