[FEATURE] [WAN2.2] Add support to train LoRA for WanVacePipeline

[SlimRG]

I can't train LoRA (aka PEFT) model, using only diffusers.
I know about DiffSync, but it has different code and some broken copabilities.

I have dataset:
In each folder 4 files:
orig.mp4 - original video (1280*720)
cropped.mp4 - cropped video with gray frames
mask.png - mask for outpaint
prompt.txt - positive prompt for outpainting

For easy navigation - I created two CSV files (train.csv and validate.csv) as in DiffSync:

video,vace_video,vace_video_mask,prompt
video-000015-1285-1-1/orig.mp4,video-000015-1285-1-1/cropped.mp4,video-000015-1285-1-1/mask.png,"Outpaint. Dim, warm ambient lighting from out-of-focus lanterns illuminates a dark interior space with indistinct wooden structures; on the left, a figure partially obscures the view, while on the right, another person moves closer before exiting frame."
...

So, I want to create pipeline:

def init_pipeline() -> WanVACEPipeline:
    """
    Init WanVACEPipeline + VAE.
    """
    vae = AutoencoderKLWan.from_pretrained(
        MODEL_ID,
        subfolder="vae",
        torch_dtype=torch.float32,
    )

    pipe: WanVACEPipeline = WanVACEPipeline.from_pretrained(
        MODEL_ID,  # "linoyts/Wan2.2-VACE-Fun-14B-diffusers"
        vae=vae,
        torch_dtype=BASE_DTYPE, # bfloat16
    )

    pipe.scheduler = UniPCMultistepScheduler.from_config(
        pipe.scheduler.config, flow_shift=FLOW_SHIFT # 5.0
    )

    # flash-attention
    if hasattr(pipe.transformer, "set_attention_backend"):
        pipe.transformer.set_attention_backend("flash")

    extra_models_to_quantize = []

    if hasattr(pipe, "transformer_2") and pipe.transformer_2 is not None:
        if hasattr(pipe.transformer_2, "set_attention_backend"):
            pipe.transformer_2.set_attention_backend("flash")
        # We train first transformer, so can quantinize second
        extra_models_to_quantize.append("transformer_2")

    # Freeze all
    if hasattr(pipe, "components"):
        for name, module in pipe.components.items():
            if isinstance(module, nn.Module):
                for p in module.parameters():
                    p.requires_grad = False
    else:
        # fallback, if no components 
        for attr in ["vae", "transformer", "transformer_2", "text_encoder"]:
            if hasattr(pipe, attr):
                m = getattr(pipe, attr)
                if isinstance(m, nn.Module):
                    for p in m.parameters():
                        p.requires_grad = False

    pipe.to(DEVICE)

    # mmgp (offload + quant)
    if ENABLE_MMGP and HAS_MMGP:
        offload.profile(
            pipe,
            MMGP_PROFILE,
            quantizeTransformer=False,        # no first transformer quant
            extraModelsToQuantize=extra_models_to_quantize,
        )
    else:
        print("[mmgp] Не используется (либо не установлен, либо отключен).")

    # Patch __call__, for remove no_grad
    _patch_pipeline_call_for_grad(pipe)

    return pipe

LoRA config:

def inject_lora_into_transformer(
    transformer: nn.Module,
    r: int,
    alpha: int,
    dropout: float,
) -> nn.Module:
    replaced = 0

    def _recursive(module: nn.Module) -> None:
        nonlocal replaced
        for name, child in list(module.named_children()):
            if isinstance(child, nn.Linear):
                setattr(module, name, LoRALinear(child, r=r, alpha=alpha, dropout=dropout))
                replaced += 1
            else:
                _recursive(child)

    _recursive(transformer)

    if replaced == 0:
        print("[LoRA] WARNING: не найдено ни одного nn.Linear в transformer для LoRA.")
    else:
        print(f"[LoRA] Обёрнуто Linear-слоёв: {replaced}")

    transformer.to(DEVICE)

    return transformer


def get_lora_parameters(transformer: nn.Module) -> List[torch.nn.Parameter]:
    params: List[torch.nn.Parameter] = []
    for name, p in transformer.named_parameters():
        if "lora_A" in name or "lora_B" in name:
            params.append(p)

    if not params:
        raise RuntimeError("Не найдено trainable LoRA-параметров (lora_A/lora_B).")

    total = sum(p.numel() for p in params)
    print(f"[LoRA] Всего LoRA-параметров: {total / 1e6:.2f}M")

    return params

When I try to train, I have error in pipe (WanVacePipeline)

Traceback (most recent call last):
  File "d:\Experiments\Train_IN2OUT\8__trainLoRA.py", line 717, in <module>
    train()
    ~~~~~^^
  File "d:\Experiments\Train_IN2OUT\8__trainLoRA.py", line 629, in train
    pipe.train()
    ^^^^^^^^^^
  File "C:\Users\BBCCA\AppData\Local\Programs\Python\Python313\Lib\site-packages\diffusers\configuration_utils.py", line 144, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'WanVACEPipeline' object has no attribute 'train'

So, can you add train() methods as in PyTorch or give an example - how to train LoRA?

---

---

---

My full code:

# -*- coding: utf-8 -*-
from __future__ import annotations

import os
import csv
import math
import random
import time
from pathlib import Path
from typing import Dict, List, Any, Tuple

import numpy as np
import av
from PIL import Image

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter

from diffusers import AutoencoderKLWan, WanVACEPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler

# mmgp (опционально)
try:
    from mmgp import offload, profile_type

    HAS_MMGP = True
except Exception:
    HAS_MMGP = False


# ==========================
#        КОНФИГ
# ==========================

# Датасет (как в твоём скрипте CSV)
DATA_ROOT = Path(r"D:/Experiments/Train_IN2OUT/4__dataset")
TRAIN_CSV_PATH = DATA_ROOT / "metadata_vace_train.csv"
VAL_CSV_PATH = DATA_ROOT / "metadata_vace_val.csv"

# Модель Wan2.2 VACE Fun
MODEL_ID = "linoyts/Wan2.2-VACE-Fun-14B-diffusers"

# Видео параметры (у тебя уже всё 1280x720x17)
FRAME_H = 720
FRAME_W = 1280
NUM_FRAMES = 17

# Обучение
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
BASE_DTYPE = torch.bfloat16  # Wan 2.x любят bfloat16
USE_AMP = True               # autocast
BATCH_SIZE = 1               # 14B, так что лучше 1
NUM_EPOCHS = 2               # подбирай под себя
MAX_TRAIN_STEPS = None       # можно оставить None, тогда считается автоматически
NUM_INFERENCE_STEPS_TRAIN = 12  # шагов диффузии в обучении (меньше ради скорости)
NUM_INFERENCE_STEPS_VAL = 16    # валидация можно чуть медленнее

LEARNING_RATE = 1e-4
WEIGHT_DECAY = 0.01

# LoRA гиперы
LORA_RANK = 32
LORA_ALPHA = 64
LORA_DROPOUT = 0.05

# Логирование и сохранение
LOG_DIR = Path(r"D:/Experiments/Train_IN2OUT/logs_wan_vace_lora")
OUTPUT_LORA_DIR = Path(r"D:/Experiments/Train_IN2OUT/wan_vace_lora_out")
SAVE_EVERY_STEPS = 2000
VAL_EVERY_STEPS = 1000
MAX_VAL_BATCHES = 64  # сколько валидационных сэмплов гонять каждый раз

# mmgp
ENABLE_MMGP = True  # ВКЛЮЧИ, если у тебя стоит mmgp и хочешь offload
MMGP_PROFILE = profile_type.HighRAM_HighVRAM if HAS_MMGP else None

# Разное
SEED = 42
FLOW_SHIFT = 5.0  # для UniPCMultistepScheduler
GUIDANCE_SCALE_TRAIN = 5.0
GUIDANCE_SCALE_VAL = 5.0

os.makedirs(LOG_DIR, exist_ok=True)
os.makedirs(OUTPUT_LORA_DIR, exist_ok=True)


# ==========================
#   Фиксация сидов
# ==========================

def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


set_seed(SEED)


# ==========================
#   Dataset / DataLoader
# ==========================

class VACELoraSample:
    """
    Одна запись датасета:
        gt_path       -> video (orig.mp4)
        cond_path     -> vace_video (cropped.mp4)
        mask_path     -> vace_video_mask (mask.png)
        prompt        -> positive prompt
    """

    __slots__ = ("gt_path", "cond_path", "mask_path", "prompt")

    def __init__(self, gt_path: str, cond_path: str, mask_path: str, prompt: str) -> None:
        self.gt_path = gt_path
        self.cond_path = cond_path
        self.mask_path = mask_path
        self.prompt = prompt


class VACELoraDataset(Dataset):
    def __init__(self, csv_path: Path, root_dir: Path) -> None:
        super().__init__()
        self.root_dir = root_dir
        self.samples: List[VACELoraSample] = []

        with csv_path.open("r", encoding="utf-8", newline="") as f:
            reader = csv.DictReader(f)
            for row in reader:
                # колонки: video, vace_video, vace_video_mask, prompt
                video_rel = row["video"]
                cond_rel = row["vace_video"]
                mask_rel = row["vace_video_mask"]
                prompt = row["prompt"]

                gt_path = (root_dir / video_rel).as_posix()
                cond_path = (root_dir / cond_rel).as_posix()
                mask_path = (root_dir / mask_rel).as_posix()

                self.samples.append(VACELoraSample(gt_path, cond_path, mask_path, prompt))

        if not self.samples:
            raise RuntimeError(f"Пустой датасет: {csv_path}")

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        s = self.samples[idx]
        # Для batch_size=1 можно не заморачиваться с коллатом
        return {
            "gt_path": s.gt_path,
            "cond_path": s.cond_path,
            "mask_path": s.mask_path,
            "prompt": s.prompt,
        }


def make_dataloaders() -> tuple[DataLoader, DataLoader]:
    train_ds = VACELoraDataset(TRAIN_CSV_PATH, DATA_ROOT)
    val_ds = VACELoraDataset(VAL_CSV_PATH, DATA_ROOT)

    train_loader = DataLoader(
        train_ds,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
        drop_last=True,
    )

    val_loader = DataLoader(
        val_ds,
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=2,
        pin_memory=True,
        drop_last=False,
    )

    return train_loader, val_loader


# ==========================
#   Видео загрузка (PyAV)
# ==========================

def load_video_pyav(
    path: str,
    num_frames: int = NUM_FRAMES,
) -> Tuple[List[Image.Image], torch.Tensor]:
    """
    Считывает видео через PyAV.
    Возвращает:
        - список PIL.Image (RGB) длиной num_frames
        - тензор torch.Tensor shape [1, num_frames, 3, H, W] в диапазоне [0,1]
    """
    container = av.open(path)
    stream = container.streams.video[0]

    frames: List[Image.Image] = []
    for frame in container.decode(stream):
        img = frame.to_rgb().to_ndarray()  # H, W, 3, uint8
        pil = Image.fromarray(img)
        frames.append(pil)
    container.close()

    if len(frames) == 0:
        raise RuntimeError(f"Видео без кадров: {path}")

    # У нас 17 кадров, но на всякий случай нормализуем количество
    if len(frames) >= num_frames:
        # Центрируем
        start = (len(frames) - num_frames) // 2
        frames = frames[start: start + num_frames]
    else:
        # Дополняем последним кадром
        last = frames[-1]
        for _ in range(num_frames - len(frames)):
            frames.append(last.copy())

    # Собираем тензор [1, T, 3, H, W] в [0,1]
    arr = np.stack(
        [np.array(f, dtype=np.float32) / 255.0 for f in frames],
        axis=0,
    )  # [T, H, W, 3]
    arr = torch.from_numpy(arr).permute(0, 3, 1, 2).unsqueeze(0)  # [1, T, 3, H, W]

    return frames, arr


# ==========================
#   LoRA-обёртка для Linear
# ==========================

class LoRALinear(nn.Module):
    """
    Простая LoRA-надстройка для nn.Linear:
        y = W x + scale * B(A x)

    Базовый Linear замораживаем, тренируем только A/B.
    """

    def __init__(
        self,
        linear: nn.Linear,
        r: int = 32,
        alpha: int = 64,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        if r <= 0:
            raise ValueError("LoRA rank r должен быть > 0")

        self.linear = linear
        self.r = r
        self.lora_A = nn.Parameter(torch.zeros(r, linear.in_features))
        self.lora_B = nn.Parameter(torch.zeros(linear.out_features, r))
        self.scaling = alpha / r
        self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()

        # Инициализация
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

        # Заморозка базового Linear
        self.linear.weight.requires_grad = False
        if self.linear.bias is not None:
            self.linear.bias.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Основной путь
        result = self.linear(x)

        # LoRA-добавка
        lora_x = self.dropout(x)
        # [*, in] -> [*, r]
        lora_x = F.linear(lora_x, self.lora_A)
        # [*, r] -> [*, out]
        lora_x = F.linear(lora_x, self.lora_B) * self.scaling

        return result + lora_x


def inject_lora_into_transformer(
    transformer: nn.Module,
    r: int,
    alpha: int,
    dropout: float,
) -> nn.Module:
    """
    Рекурсивно оборачиваем ВСЕ nn.Linear в transformer в LoRALinear.
    """
    replaced = 0

    def _recursive(module: nn.Module) -> None:
        nonlocal replaced
        for name, child in list(module.named_children()):
            if isinstance(child, nn.Linear):
                setattr(module, name, LoRALinear(child, r=r, alpha=alpha, dropout=dropout))
                replaced += 1
            else:
                _recursive(child)

    _recursive(transformer)

    if replaced == 0:
        print("[LoRA] WARNING: не найдено ни одного nn.Linear в transformer для LoRA.")
    else:
        print(f"[LoRA] Обёрнуто Linear-слоёв: {replaced}")

    # Переносим transformer (включая LoRA-параметры) на нужное устройство
    transformer.to(DEVICE)

    return transformer


def get_lora_parameters(transformer: nn.Module) -> List[torch.nn.Parameter]:
    """
    Собираем только LoRA-параметры (для оптимизатора).
    """
    params: List[torch.nn.Parameter] = []
    for name, p in transformer.named_parameters():
        if "lora_A" in name or "lora_B" in name:
            params.append(p)

    if not params:
        raise RuntimeError("Не найдено trainable LoRA-параметров (lora_A/lora_B).")

    total = sum(p.numel() for p in params)
    print(f"[LoRA] Всего LoRA-параметров: {total / 1e6:.2f}M")

    return params


def save_lora_checkpoint(
    pipe: WanVACEPipeline,
    step: int,
    output_dir: Path,
) -> None:
    """
    Сохранение LoRA-адаптера:
      - сохраняем только state_dict LoRALinear-слоёв transformer.
    Позже можно будет написать маленький скрипт-конвертер в любой удобный формат.
    """
    ckpt_dir = output_dir / f"step_{step:07d}"
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # Собираем только LoRA-параметры transformer'а
    lora_state = {
        k: v.cpu()
        for k, v in pipe.transformer.state_dict().items()
        if "lora_A" in k or "lora_B" in k
    }

    ckpt_path = ckpt_dir / "lora_transformer.pt"
    print(f"[SAVE] Сохраняем LoRA в {ckpt_path} ...")
    torch.save(lora_state, ckpt_path)
    print("[SAVE] Готово.")


# ==========================
#   Инициализация пайплайна
# ==========================

def _patch_pipeline_call_for_grad(pipe: WanVACEPipeline) -> None:
    """
    Diffusers-пайплайны часто декорированы @torch.no_grad().
    Здесь аккуратно вытаскиваем оригинальный __call__ без no_grad, если возможно.
    """
    call_fn = pipe.__call__
    orig_call = getattr(call_fn, "__wrapped__", None)

    if orig_call is not None:
        # привязываем как bound-method к объекту pipe
        pipe._call_with_grad = orig_call.__get__(pipe, type(pipe))
        print("[PIPE] Найден __wrapped__; будем использовать pipe._call_with_grad() без no_grad.")
    else:
        # fallback — будет использоваться как есть (если вдруг без декоратора)
        pipe._call_with_grad = pipe.__call__
        print("[PIPE] __wrapped__ не найден; используем pipe.__call__ как есть.")


def init_pipeline() -> WanVACEPipeline:
    """
    Инициализация WanVACEPipeline + VAE.
    """
    print("Загрузка VAE (bfloat16)...")
    vae = AutoencoderKLWan.from_pretrained(
        MODEL_ID,
        subfolder="vae",
        torch_dtype=torch.bfloat16,
    )

    print("Загрузка WanVACEPipeline (bfloat16)...")
    pipe: WanVACEPipeline = WanVACEPipeline.from_pretrained(
        MODEL_ID,
        vae=vae,
        torch_dtype=BASE_DTYPE,
    )

    # UniPC с flow_shift как ты обычно делаешь
    pipe.scheduler = UniPCMultistepScheduler.from_config(
        pipe.scheduler.config, flow_shift=FLOW_SHIFT
    )

    # Включаем flash-attention, если поддерживается
    if hasattr(pipe.transformer, "set_attention_backend"):
        pipe.transformer.set_attention_backend("flash")

    extra_models_to_quantize = []

    if hasattr(pipe, "transformer_2") and pipe.transformer_2 is not None:
        if hasattr(pipe.transformer_2, "set_attention_backend"):
            pipe.transformer_2.set_attention_backend("flash")
        # второму трансформеру LoRA не навешиваем → его можно квантизовать/offload'ить
        extra_models_to_quantize.append("transformer_2")

    # Замораживаем ВСЕ модули пайплайна (vae, transformer, transformer_2, text_encoder, ...)
    if hasattr(pipe, "components"):
        for name, module in pipe.components.items():
            if isinstance(module, nn.Module):
                for p in module.parameters():
                    p.requires_grad = False
    else:
        # fallback, если components нет (на всякий случай)
        for attr in ["vae", "transformer", "transformer_2", "text_encoder"]:
            if hasattr(pipe, attr):
                m = getattr(pipe, attr)
                if isinstance(m, nn.Module):
                    for p in m.parameters():
                        p.requires_grad = False

    # Перенос на устройство (базовые веса)
    pipe.to(DEVICE)

    # mmgp (offload + quant только для второго трансформера)
    if ENABLE_MMGP and HAS_MMGP:
        print("[mmgp] Включаем offload.profile (quantize только transformer_2).")
        offload.profile(
            pipe,
            MMGP_PROFILE,
            quantizeTransformer=False,        # первый transformer НЕ квантуем
            extraModelsToQuantize=extra_models_to_quantize,
        )
    else:
        print("[mmgp] Не используется (либо не установлен, либо отключен).")

    # Патчим __call__, чтобы убрать no_grad
    _patch_pipeline_call_for_grad(pipe)

    return pipe


# ==========================
#   Вспомогательная генерация
# ==========================

def run_vace_forward(
    pipe: WanVACEPipeline,
    cond_frames: List[Image.Image],
    mask_img: Image.Image,
    prompt: str,
    num_inference_steps: int,
    guidance_scale: float,
) -> torch.Tensor:
    """
    Один прогон WanVACEPipeline.
    Возврат: тензор видео [1, T, 3, H, W] в диапазоне [0,1] (output_type="pt").
    ВАЖНО: используем pipe._call_with_grad, чтобы autograd работал.
    """
    # Маска одна на всё видео → размножаем по кадрам
    mask_list = [mask_img] * len(cond_frames)

    out = pipe._call_with_grad(
        video=cond_frames,
        mask=mask_list,
        prompt=prompt,
        height=FRAME_H,
        width=FRAME_W,
        num_frames=len(cond_frames),
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        output_type="pt",   # критично для градиента
        return_dict=True,
    )

    frames = out.frames  # обычно torch.Tensor или np.ndarray

    if isinstance(frames, torch.Tensor):
        video = frames  # ожидаем [B, T, C, H, W] или [T, C, H, W]
    else:
        # np.ndarray или список → в тензор
        frames_np = np.array(frames)
        video = torch.from_numpy(frames_np)

    if video.dim() == 4:
        # [T, C, H, W] → [1, T, C, H, W]
        video = video.unsqueeze(0)

    # Приводим к [1, T, 3, H, W]
    if video.shape[2] != 3:
        raise RuntimeError(f"Ожидали 3 канала, получили: {video.shape}")

    # Диапазон diffusers обычно [0,1] для pt-выхода; оставляем как есть.
    return video


# ==========================
#     Train / Val циклы
# ==========================

def validate(
    pipe: WanVACEPipeline,
    val_loader: DataLoader,
    writer: SummaryWriter,
    global_step: int,
) -> float:
    pipe.eval()
    total_loss = 0.0
    count = 0

    with torch.no_grad():  # валидация без градиента
        for batch_idx, batch in enumerate(val_loader):
            if batch_idx >= MAX_VAL_BATCHES:
                break

            gt_path = batch["gt_path"]
            cond_path = batch["cond_path"]
            mask_path = batch["mask_path"]
            prompt = batch["prompt"]

            # Для batch_size=1 превращаем в строки
            if isinstance(gt_path, list):
                gt_path = gt_path[0]
                cond_path = cond_path[0]
                mask_path = mask_path[0]
                prompt = prompt[0]

            # Загружаем GT и condition
            cond_frames, _ = load_video_pyav(cond_path, num_frames=NUM_FRAMES)
            _, gt_tensor = load_video_pyav(gt_path, num_frames=NUM_FRAMES)
            gt_tensor = gt_tensor.to(DEVICE, dtype=torch.float32)

            mask_img = Image.open(mask_path).convert("L")

            with torch.cuda.amp.autocast(enabled=USE_AMP, dtype=BASE_DTYPE):
                pred_video = run_vace_forward(
                    pipe,
                    cond_frames,
                    mask_img,
                    prompt,
                    num_inference_steps=NUM_INFERENCE_STEPS_VAL,
                    guidance_scale=GUIDANCE_SCALE_VAL,
                )
                pred_video = pred_video.to(DEVICE, dtype=torch.float32)

                # Выравниваем T
                min_frames = min(pred_video.shape[1], gt_tensor.shape[1])
                pred_use = pred_video[:, :min_frames]
                gt_use = gt_tensor[:, :min_frames]

                loss = F.l1_loss(pred_use, gt_use)

            total_loss += loss.item()
            count += 1

            # Для первых пар трекаем картинку в TensorBoard
            if batch_idx == 0:
                # frame 0 (gt vs pred)
                gt_frame = gt_use[0, 0]  # [3,H,W]
                pred_frame = pred_use[0, 0]
                # Собираем грид 1x2
                grid = torch.stack([gt_frame, pred_frame], dim=0)
                writer.add_images("val/gt_pred_frame0", grid, global_step)

    avg_loss = total_loss / max(1, count)
    writer.add_scalar("val/loss", avg_loss, global_step)
    print(f"[VAL] step={global_step} loss={avg_loss:.6f} (по {count} батчам)")
    return avg_loss


def train() -> None:
    train_loader, val_loader = make_dataloaders()

    pipe = init_pipeline()

    # Вешаем LoRA на первый transformer
    print("[LoRA] Инжектим LoRA в pipe.transformer ...")
    pipe.transformer = inject_lora_into_transformer(
        pipe.transformer,
        r=LORA_RANK,
        alpha=LORA_ALPHA,
        dropout=LORA_DROPOUT,
    )

    # Собираем trainable LoRA-параметры
    lora_params = get_lora_parameters(pipe.transformer)

    optimizer = torch.optim.AdamW(
        lora_params,
        lr=LEARNING_RATE,
        weight_decay=WEIGHT_DECAY,
        betas=(0.9, 0.999),
        eps=1e-8,
    )

    writer = SummaryWriter(LOG_DIR)

    # Подсчёт максимального количества шагов, если не задано
    if MAX_TRAIN_STEPS is not None:
        max_steps = MAX_TRAIN_STEPS
    else:
        max_steps = NUM_EPOCHS * math.ceil(len(train_loader.dataset) / BATCH_SIZE)

    print(f"Старт обучения: max_steps={max_steps}, epochs={NUM_EPOCHS}")

    global_step = 0
    start_time = time.time()

    scaler = torch.cuda.amp.GradScaler(enabled=USE_AMP)

    for epoch in range(NUM_EPOCHS):
        pipe.train()
        epoch_loss = 0.0
        for batch_idx, batch in enumerate(train_loader):
            if global_step >= max_steps:
                break

            gt_path = batch["gt_path"]
            cond_path = batch["cond_path"]
            mask_path = batch["mask_path"]
            prompt = batch["prompt"]

            # batch_size=1 → распаковываем
            if isinstance(gt_path, list):
                gt_path = gt_path[0]
                cond_path = cond_path[0]
                mask_path = mask_path[0]
                prompt = prompt[0]

            cond_frames, _ = load_video_pyav(cond_path, num_frames=NUM_FRAMES)
            _, gt_tensor = load_video_pyav(gt_path, num_frames=NUM_FRAMES)
            gt_tensor = gt_tensor.to(DEVICE, dtype=torch.float32)

            mask_img = Image.open(mask_path).convert("L")

            optimizer.zero_grad(set_to_none=True)

            with torch.cuda.amp.autocast(enabled=USE_AMP, dtype=BASE_DTYPE):
                pred_video = run_vace_forward(
                    pipe,
                    cond_frames,
                    mask_img,
                    prompt,
                    num_inference_steps=NUM_INFERENCE_STEPS_TRAIN,
                    guidance_scale=GUIDANCE_SCALE_TRAIN,
                )
                pred_video = pred_video.to(DEVICE, dtype=torch.float32)

                # Выравниваем по числу кадров
                min_frames = min(pred_video.shape[1], gt_tensor.shape[1])
                pred_use = pred_video[:, :min_frames]
                gt_use = gt_tensor[:, :min_frames]

                loss = F.l1_loss(pred_use, gt_use)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            global_step += 1
            epoch_loss += loss.item()

            if global_step % 10 == 0:
                avg_loss = epoch_loss / (batch_idx + 1)
                elapsed = time.time() - start_time
                print(
                    f"[TRAIN] epoch={epoch+1} step={global_step}/{max_steps} "
                    f"batch_loss={loss.item():.6f} avg_epoch_loss={avg_loss:.6f} "
                    f"time={elapsed/60:.1f}min"
                )
                writer.add_scalar("train/loss", loss.item(), global_step)

            # Периодическая валидация
            if global_step % VAL_EVERY_STEPS == 0:
                _ = validate(pipe, val_loader, writer, global_step)
                pipe.train()

            # Периодическое сохранение LoRA
            if global_step % SAVE_EVERY_STEPS == 0:
                save_lora_checkpoint(pipe, global_step, OUTPUT_LORA_DIR)

            if global_step >= max_steps:
                break

        # конец эпохи
        avg_epoch_loss = epoch_loss / max(1, (batch_idx + 1))
        writer.add_scalar("train/epoch_loss", avg_epoch_loss, epoch + 1)
        print(f"=== Epoch {epoch+1}/{NUM_EPOCHS} завершена. avg_loss={avg_epoch_loss:.6f} ===")

        # Валидация в конце эпохи
        _ = validate(pipe, val_loader, writer, global_step)

    # Финальное сохранение
    save_lora_checkpoint(pipe, global_step, OUTPUT_LORA_DIR)
    writer.close()
    print("Обучение завершено.")


if __name__ == "__main__":
    train()

@DN6 @a-r-r-o-w

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.