From 10e50d5797206644376bb65c5a7ec0b49fbd3bd0 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 12 Jun 2026 12:16:51 -0600 Subject: [PATCH] Fix issue with casting unet after compilation --- jobs/process/BaseSDTrainProcess.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 1a894aa..ca43531 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -2215,6 +2215,14 @@ class BaseSDTrainProcess(BaseTrainProcess): fullgraph=compile_fullgraph, ) + if not is_unet_offloaded: + # once compiled, dynamo guards hold weakrefs to the params; + # .to() on quantized params requires swap_tensors, which fails + # on tensors with weakrefs. The model stays on device anyway, + # so make .to() a no-op. + unet_module = self.sd.unet + unet_module.to = lambda *args, **kwargs: unet_module + except Exception as e: print_acc(f"Failed to compile model: {e}") print_acc("Continuing without compilation")