diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 1624fc4f..475176ad 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -2049,8 +2049,10 @@ class BaseSDTrainProcess(BaseTrainProcess): # compile the model if needed (must be after LoRA/adapter injection AND accelerator.prepare) if self.model_config.compile: try: - print_acc(f"Compiling model with torch.compile") - self.sd.unet = torch.compile(self.sd.unet, dynamic=True, mode='reduce-overhead') + # make sure it is on the gpu + self.sd.unet.to(self.device_torch) + print_acc("Compiling model with torch.compile. The first forward will hang for a while using this. This is normal.") + self.sd.unet = torch.compile(self.sd.unet) except Exception as e: print_acc(f"Failed to compile model: {e}") print_acc("Continuing without compilation") diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 03d22306..09315aed 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -695,6 +695,10 @@ class ModelConfig: # compile the model with torch compile self.compile = kwargs.get("compile", False) + if self.compile and self.quantize: + print("Warning: You cannot compile a quantized model. Disabling compile.") + self.compile = False + # kwargs to pass to the model self.model_kwargs = kwargs.get("model_kwargs", {})