diffsynth/diffusion/runner.py

import hashlib
import json
import os, torch
from tqdm import tqdm
from accelerate import Accelerator
from .training_module import DiffusionTrainingModule
from .logger import ModelLogger

RESUME_META_FILE = "resume_meta.json"


def _checkpoint_dir(output_path, step):
    return os.path.join(output_path, f"training_state-step-{step}")


def _find_latest_checkpoint_dir(output_path):
    if not os.path.isdir(output_path):
        return None
    candidates = []
    for name in os.listdir(output_path):
        if not name.startswith("training_state-step-"):
            continue
        try:
            step = int(name.rsplit("-", 1)[-1])
        except ValueError:
            continue
        path = os.path.join(output_path, name)
        if os.path.isdir(path):
            candidates.append((step, path))
    if not candidates:
        return None
    candidates.sort(key=lambda item: item[0])
    return candidates[-1][1]


def _save_resume_state(accelerator, output_path, global_step, epoch_id, step_in_epoch):
    accelerator.wait_for_everyone()
    checkpoint_path = _checkpoint_dir(output_path, global_step)
    accelerator.save_state(checkpoint_path)
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        meta = {
            "global_step": int(global_step),
            "epoch_id": int(epoch_id),
            "step_in_epoch": int(step_in_epoch),
        }
        with open(os.path.join(checkpoint_path, RESUME_META_FILE), "w", encoding="utf-8") as handle:
            json.dump(meta, handle)


def _load_resume_state(accelerator, resume_from, output_path):
    if resume_from is None:
        return None
    if resume_from == "latest":
        checkpoint_path = _find_latest_checkpoint_dir(output_path)
        if checkpoint_path is None:
            return None
    else:
        checkpoint_path = os.path.abspath(resume_from)
        if not os.path.isdir(checkpoint_path):
            raise FileNotFoundError(f"Resume checkpoint folder does not exist: {checkpoint_path}")

    accelerator.load_state(checkpoint_path)
    meta_path = os.path.join(checkpoint_path, RESUME_META_FILE)
    if os.path.isfile(meta_path):
        with open(meta_path, "r", encoding="utf-8") as handle:
            meta = json.load(handle)
    else:
        meta = {"global_step": 0, "epoch_id": 0, "step_in_epoch": 0}
    meta["checkpoint_path"] = checkpoint_path
    return meta


def _module_resume_fingerprint(unwrapped_model, module_attr="dit", max_tensors=8, sample_size=2048):
    pipe = getattr(unwrapped_model, "pipe", None)
    module = getattr(pipe, module_attr, None) if pipe is not None else None
    if module is None:
        return None

    entries = []
    for name, tensor in sorted(module.state_dict().items()):
        if not torch.is_tensor(tensor):
            continue
        if not tensor.is_floating_point():
            continue
        if tensor.numel() == 0:
            continue
        flat = tensor.detach().float().view(-1)
        n = min(sample_size, flat.numel())
        sample = flat[:n].cpu()
        sample_bytes = sample.numpy().tobytes()
        entries.append({
            "name": name,
            "shape": tuple(tensor.shape),
            "hash": hashlib.sha1(sample_bytes).hexdigest()[:12],
            "mean": float(sample.mean().item()),
            "std": float(sample.std(unbiased=False).item()) if n > 1 else 0.0,
            "norm": float(sample.norm().item()),
        })
        if len(entries) >= max_tensors:
            break
    return {"module": f"pipe.{module_attr}", "entries": entries}


def _compare_fingerprints(before_fp, after_fp):
    if before_fp is None or after_fp is None:
        return None
    before = {entry["name"]: entry["hash"] for entry in before_fp.get("entries", [])}
    after = {entry["name"]: entry["hash"] for entry in after_fp.get("entries", [])}
    common = sorted(set(before.keys()) & set(after.keys()))
    changed = [name for name in common if before[name] != after[name]]
    return {
        "total_common": len(common),
        "changed": changed,
    }


def launch_training_task(
    accelerator: Accelerator,
    dataset: torch.utils.data.Dataset,
    model: DiffusionTrainingModule,
    model_logger: ModelLogger,
    learning_rate: float = 1e-5,
    weight_decay: float = 1e-2,
    num_workers: int = 1,
    save_steps: int = None,
    num_epochs: int = 1,
    args = None,
):
    if args is not None:
        learning_rate = args.learning_rate
        weight_decay = args.weight_decay
        num_workers = args.dataset_num_workers
        save_steps = args.save_steps
        num_epochs = args.num_epochs
    
    optimizer = torch.optim.AdamW(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay)
    scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
    dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=num_workers)
    
    model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
    
    resume_from = getattr(args, "resume_from", None) if args is not None else None
    pre_resume_fp = None
    if resume_from is not None:
        unwrapped_model = accelerator.unwrap_model(model)
        pre_resume_fp = _module_resume_fingerprint(unwrapped_model)
        if accelerator.is_main_process and pre_resume_fp is not None:
            print(
                f"[ResumeCheck] Pre-resume fingerprint on {pre_resume_fp['module']} "
                f"(sample_tensors={len(pre_resume_fp['entries'])})"
            )

    resume_meta = _load_resume_state(
        accelerator=accelerator,
        resume_from=resume_from,
        output_path=getattr(args, "output_path", "./models") if args is not None else "./models",
    )
    start_epoch = 0
    start_step_in_epoch = 0
    global_step = 0
    if resume_meta is not None:
        global_step = int(resume_meta.get("global_step", 0))
        start_epoch = int(resume_meta.get("epoch_id", 0))
        start_step_in_epoch = int(resume_meta.get("step_in_epoch", 0))
        model_logger.num_steps = global_step
        unwrapped_model = accelerator.unwrap_model(model)
        if hasattr(unwrapped_model, "_debug_forward_step"):
            unwrapped_model._debug_forward_step = global_step
        if hasattr(unwrapped_model, "pipe"):
            unwrapped_model.pipe._debug_forward_step = global_step
        if accelerator.is_main_process:
            print(f"[Resume] Loaded state from {resume_meta.get('checkpoint_path')}")
            print(f"[Resume] global_step={global_step}, epoch={start_epoch}, step_in_epoch={start_step_in_epoch}")
            post_resume_fp = _module_resume_fingerprint(unwrapped_model)
            fp_cmp = _compare_fingerprints(pre_resume_fp, post_resume_fp)
            if fp_cmp is not None:
                print(
                    f"[ResumeCheck] Fingerprint changed tensors: "
                    f"{len(fp_cmp['changed'])}/{fp_cmp['total_common']}"
                )
                if len(fp_cmp["changed"]) == 0:
                    print("[ResumeCheck][WARN] No sampled tensor changed after resume load.")
                else:
                    print(f"[ResumeCheck] Changed (first 5): {fp_cmp['changed'][:5]}")
    elif resume_from is not None and accelerator.is_main_process:
        print(f"[Resume] No checkpoint found for resume_from={resume_from}. Starting from scratch.")

    for epoch_id in range(start_epoch, num_epochs):
        progress_bar = tqdm(
            dataloader, 
            desc=f"Epoch {epoch_id + 1}/{num_epochs}", 
            disable=not accelerator.is_main_process,
            leave=True
        )
        for step_in_epoch, data in enumerate(progress_bar):
            if epoch_id == start_epoch and step_in_epoch < start_step_in_epoch:
                continue
            # Skip empty or invalid samples.
            if data is None:
                continue
            
            # Extra guard: skip if the parsed sample has no video payload.
            if data.get("video") is None:
                continue
            with accelerator.accumulate(model):
                optimizer.zero_grad()
                if dataset.load_from_cache:
                    loss = model({}, inputs=data)
                else:
                    loss = model(data)
                accelerator.backward(loss)
                optimizer.step()
                global_step += 1
                model_logger.on_step_end(accelerator, model, save_steps)
                scheduler.step()
                if save_steps is not None and global_step % save_steps == 0:
                    _save_resume_state(
                        accelerator=accelerator,
                        output_path=model_logger.output_path,
                        global_step=global_step,
                        epoch_id=epoch_id,
                        step_in_epoch=step_in_epoch + 1,
                    )
                if accelerator.is_main_process:
                    # Read current learning rate.
                    current_lr = 0.0
                    if hasattr(optimizer, "param_groups"):
                        current_lr = optimizer.param_groups[0]['lr']
                    
                    # Refresh progress bar display.
                    progress_bar.set_postfix(
                        loss=f"{loss.item():.4f}", 
                        lr=f"{current_lr:.1e}"
                    )
        start_step_in_epoch = 0
        if save_steps is None:
            model_logger.on_epoch_end(accelerator, model, epoch_id)
            _save_resume_state(
                accelerator=accelerator,
                output_path=model_logger.output_path,
                global_step=global_step,
                epoch_id=epoch_id + 1,
                step_in_epoch=0,
            )
    model_logger.on_training_end(accelerator, model, save_steps)


def launch_data_process_task(
    accelerator: Accelerator,
    dataset: torch.utils.data.Dataset,
    model: DiffusionTrainingModule,
    model_logger: ModelLogger,
    num_workers: int = 8,
    args = None,
):
    if args is not None:
        num_workers = args.dataset_num_workers
        
    dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, collate_fn=lambda x: x[0], num_workers=num_workers)
    model, dataloader = accelerator.prepare(model, dataloader)
    
    for data_id, data in enumerate(tqdm(dataloader)):
        with accelerator.accumulate(model):
            with torch.no_grad():
                folder = os.path.join(model_logger.output_path, str(accelerator.process_index))
                os.makedirs(folder, exist_ok=True)
                save_path = os.path.join(model_logger.output_path, str(accelerator.process_index), f"{data_id}.pth")
                data = model(data)
                torch.save(data, save_path)