xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-03 21:17:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use encoded latents as input

Browse Source
This commit is contained in:
Kohaku-Blueleaf 2025-06-13 11:49:49 +08:00
parent 23523f5894
commit 218c3e3793
1 changed files with 9 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
comfy_extras/nodes_train.py
View File

@ -262,20 +262,16 @@ class TrainLoraNode:
return {
"required": {
"model": (IO.MODEL, {"tooltip": "The model to train the LoRA on."}),
"vae": (
IO.VAE,
"latents": (
"LATENT",
{
"tooltip": "The VAE model to use for encoding images for training."
"tooltip": "The Latents to use for training, serve as dataset/input of the model."
},
),
"positive": (
IO.CONDITIONING,
{"tooltip": "The positive conditioning to use for training."},
),
"image": (
IO.IMAGE,
{"tooltip": "The image or image batch to train the LoRA on."},
),
"batch_size": (
IO.INT,
{
@ -364,9 +360,8 @@ class TrainLoraNode:
def train(
self,
model,
vae,
latents,
positive,
image,
batch_size,
steps,
learning_rate,
@ -378,17 +373,14 @@ class TrainLoraNode:
lora_dtype,
existing_lora,
):
num_images = image.shape[0]
indices = torch.randperm(num_images)[:batch_size]
batch_tensor = image[indices]
# Ensure we're not in inference mode when encoding
encoded = vae.encode(batch_tensor)
mp = model.clone()
dtype = node_helpers.string_to_torch_dtype(training_dtype)
lora_dtype = node_helpers.string_to_torch_dtype(lora_dtype)
mp.set_model_compute_dtype(dtype)
latents = latents["samples"].to(dtype)
num_images = latents.shape[0]
with torch.inference_mode(False):
lora_sd = {}
generator = torch.Generator()
@ -509,8 +501,9 @@ class TrainLoraNode:
noise = comfy_extras.nodes_custom_sampler.Noise_RandomNoise(step * 1000 + seed)
indices = torch.randperm(num_images)[:batch_size]
ss.sample(
noise, guider, train_sampler, sigma, {"samples": encoded.clone()}
noise, guider, train_sampler, sigma, {"samples": latents[indices].clone()}
)
finally:
for m in mp.model.modules():