xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-03 22:27:13 +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 { return {
"required": { "required": {
"model": (IO.MODEL, {"tooltip": "The model to train the LoRA on."}), "model": (IO.MODEL, {"tooltip": "The model to train the LoRA on."}),
"vae": ( "latents": (
IO.VAE, "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": ( "positive": (
IO.CONDITIONING, IO.CONDITIONING,
{"tooltip": "The positive conditioning to use for training."}, {"tooltip": "The positive conditioning to use for training."},
), ),
"image": (
IO.IMAGE,
{"tooltip": "The image or image batch to train the LoRA on."},
),
"batch_size": ( "batch_size": (
IO.INT, IO.INT,
{ {
@ -364,9 +360,8 @@ class TrainLoraNode:
def train( def train(
self, self,
model, model,
vae, latents,
positive, positive,
image,
batch_size, batch_size,
steps, steps,
learning_rate, learning_rate,
@ -378,17 +373,14 @@ class TrainLoraNode:
lora_dtype, lora_dtype,
existing_lora, 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() mp = model.clone()
dtype = node_helpers.string_to_torch_dtype(training_dtype) dtype = node_helpers.string_to_torch_dtype(training_dtype)
lora_dtype = node_helpers.string_to_torch_dtype(lora_dtype) lora_dtype = node_helpers.string_to_torch_dtype(lora_dtype)
mp.set_model_compute_dtype(dtype) mp.set_model_compute_dtype(dtype)
latents = latents["samples"].to(dtype)
num_images = latents.shape[0]
with torch.inference_mode(False): with torch.inference_mode(False):
lora_sd = {} lora_sd = {}
generator = torch.Generator() generator = torch.Generator()
@ -509,8 +501,9 @@ class TrainLoraNode:
noise = comfy_extras.nodes_custom_sampler.Noise_RandomNoise(step * 1000 + seed) noise = comfy_extras.nodes_custom_sampler.Noise_RandomNoise(step * 1000 + seed)
indices = torch.randperm(num_images)[:batch_size]
ss.sample( ss.sample(
noise, guider, train_sampler, sigma, {"samples": encoded.clone()} noise, guider, train_sampler, sigma, {"samples": latents[indices].clone()}
) )
finally: finally:
for m in mp.model.modules(): for m in mp.model.modules():