[Trainer/bug] Ensure model is not inference mode (CORE-72) (#13400)

* Ensure model is not inference mode

* force clone inside training mode to avoid inference tensor

* Allow force deepcopy for model patcher
This commit is contained in:
Kohaku-Blueleaf 2026-06-10 11:07:47 +08:00 committed by GitHub
parent 46d45aade1
commit f350acdf21
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 41 additions and 40 deletions

View File

@ -379,8 +379,9 @@ class ModelPatcher:
def get_clone_model_override(self):
return self.model, (self.backup, self.backup_buffers, self.object_patches_backup, self.pinned)
def clone(self, disable_dynamic=False, model_override=None):
def clone(self, disable_dynamic=False, model_override=None, force_deepcopy=False):
class_ = self.__class__
if self.is_dynamic() and disable_dynamic or force_deepcopy:
if self.is_dynamic() and disable_dynamic:
class_ = ModelPatcher
if model_override is None:

View File

@ -1149,8 +1149,9 @@ class TrainLoraNode(io.ComfyNode):
# Process conditioning
positive = _process_conditioning(positive)
with torch.inference_mode(False):
# Setup model and dtype
mp = model.clone()
mp = model.clone(force_deepcopy=True)
use_grad_scaler = False
lora_dtype = node_helpers.string_to_torch_dtype(lora_dtype)
if training_dtype != "none":
@ -1185,9 +1186,8 @@ class TrainLoraNode(io.ComfyNode):
# Validate and expand conditioning
positive = _validate_and_expand_conditioning(positive, num_images, bucket_mode)
with torch.inference_mode(False):
# Setup models for training
mp.model.requires_grad_(False)
mp.model.requires_grad_(False).train()
# Load existing LoRA weights if provided
existing_weights, existing_steps = _load_existing_lora(existing_lora)