(self, loadmodel_names=[])
| 93 | |
| 94 | |
| 95 | def load_models_to_device(self, loadmodel_names=[]): |
| 96 | # only load models to device if cpu_offload is enabled |
| 97 | if not self.cpu_offload: |
| 98 | return |
| 99 | # offload the unneeded models to cpu |
| 100 | for model_name in self.model_names: |
| 101 | if model_name not in loadmodel_names: |
| 102 | model = getattr(self, model_name) |
| 103 | if model is not None: |
| 104 | if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: |
| 105 | for module in model.modules(): |
| 106 | if hasattr(module, "offload"): |
| 107 | module.offload() |
| 108 | else: |
| 109 | model.cpu() |
| 110 | # load the needed models to device |
| 111 | for model_name in loadmodel_names: |
| 112 | model = getattr(self, model_name) |
| 113 | if model is not None: |
| 114 | if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: |
| 115 | for module in model.modules(): |
| 116 | if hasattr(module, "onload"): |
| 117 | module.onload() |
| 118 | else: |
| 119 | model.to(self.device) |
| 120 | # fresh the cuda cache |
| 121 | torch.cuda.empty_cache() |
| 122 | |
| 123 | |
| 124 | def generate_noise(self, shape, seed=None, device="cpu", dtype=torch.float16): |
no test coverage detected