| 36 | |
| 37 | |
| 38 | def enable_vram_management(self, num_persistent_param_in_dit=None): |
| 39 | dtype = next(iter(self.text_encoder.parameters())).dtype |
| 40 | enable_vram_management( |
| 41 | self.text_encoder, |
| 42 | module_map = { |
| 43 | torch.nn.Linear: AutoWrappedLinear, |
| 44 | torch.nn.Embedding: AutoWrappedModule, |
| 45 | T5RelativeEmbedding: AutoWrappedModule, |
| 46 | T5LayerNorm: AutoWrappedModule, |
| 47 | }, |
| 48 | module_config = dict( |
| 49 | offload_dtype=dtype, |
| 50 | offload_device="cpu", |
| 51 | onload_dtype=dtype, |
| 52 | onload_device="cpu", |
| 53 | computation_dtype=self.torch_dtype, |
| 54 | computation_device=self.device, |
| 55 | ), |
| 56 | ) |
| 57 | dtype = next(iter(self.dit.parameters())).dtype |
| 58 | enable_vram_management( |
| 59 | self.dit, |
| 60 | module_map = { |
| 61 | torch.nn.Linear: AutoWrappedLinear, |
| 62 | torch.nn.Conv3d: AutoWrappedModule, |
| 63 | torch.nn.LayerNorm: AutoWrappedModule, |
| 64 | RMSNorm: AutoWrappedModule, |
| 65 | }, |
| 66 | module_config = dict( |
| 67 | offload_dtype=dtype, |
| 68 | offload_device="cpu", |
| 69 | onload_dtype=dtype, |
| 70 | onload_device=self.device, |
| 71 | computation_dtype=self.torch_dtype, |
| 72 | computation_device=self.device, |
| 73 | ), |
| 74 | max_num_param=num_persistent_param_in_dit, |
| 75 | overflow_module_config = dict( |
| 76 | offload_dtype=dtype, |
| 77 | offload_device="cpu", |
| 78 | onload_dtype=dtype, |
| 79 | onload_device="cpu", |
| 80 | computation_dtype=self.torch_dtype, |
| 81 | computation_device=self.device, |
| 82 | ), |
| 83 | ) |
| 84 | dtype = next(iter(self.vae.parameters())).dtype |
| 85 | enable_vram_management( |
| 86 | self.vae, |
| 87 | module_map = { |
| 88 | torch.nn.Linear: AutoWrappedLinear, |
| 89 | torch.nn.Conv2d: AutoWrappedModule, |
| 90 | RMS_norm: AutoWrappedModule, |
| 91 | CausalConv3d: AutoWrappedModule, |
| 92 | Upsample: AutoWrappedModule, |
| 93 | torch.nn.SiLU: AutoWrappedModule, |
| 94 | torch.nn.Dropout: AutoWrappedModule, |
| 95 | }, |