train/train_kontext_lora.py

import argparse
import copy
import logging
import math
import os
import shutil
from contextlib import nullcontext
from pathlib import Path
import re
import time

from safetensors.torch import save_file
from PIL import Image
import numpy as np
import torch
import torch.utils.checkpoint
import transformers

from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed

import diffusers
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxPipeline
from diffusers.optimization import get_scheduler
from diffusers.training_utils import (
    cast_training_params,
    compute_density_for_timestep_sampling,
    compute_loss_weighting_for_sd3,
)
from diffusers.utils.torch_utils import is_compiled_module
from diffusers.utils import (
    check_min_version,
    is_wandb_available,
)

from src.prompt_helper import *
from src.lora_helper import *
from src.jsonl_datasets_kontext_interactive_lora import make_interactive_dataset_subjects, make_placement_dataset_subjects, make_pexels_dataset_subjects, make_mixed_dataset, collate_fn
from diffusers import FluxKontextPipeline
from diffusers.models import FluxTransformer2DModel
from tqdm.auto import tqdm
from peft import LoraConfig
from peft.utils import get_peft_model_state_dict
from diffusers.utils import convert_state_dict_to_diffusers

if is_wandb_available():
    import wandb


# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
check_min_version("0.31.0.dev0")

logger = get_logger(__name__)


PREFERRED_KONTEXT_RESOLUTIONS = [
    (672, 1568),
    (688, 1504),
    (720, 1456),
    (752, 1392),
    (832, 1248),
    (880, 1184),
    (944, 1104),
    (1024, 1024),
    (1104, 944),
    (1184, 880),
    (1248, 832),
    (1392, 752),
    (1456, 720),
    (1504, 688),
    (1568, 672),
]


def log_validation(
    pipeline,
    args,
    accelerator,
    pipeline_args,
    step,
    torch_dtype,
    is_final_validation=False,
):
    logger.info(
        f"Running validation... Paired evaluation for image and prompt."
    )
    pipeline = pipeline.to(device=accelerator.device, dtype=torch_dtype)
    pipeline.set_progress_bar_config(disable=True)

    generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
    # Match compute dtype for validation to avoid dtype mismatches (e.g., VAE bf16 vs float latents)
    if torch_dtype in (torch.float16, torch.bfloat16):
        device_type = 'cuda' if torch.cuda.is_available() else 'cpu'
        autocast_ctx = torch.autocast(device_type=device_type, dtype=torch_dtype)
    else:
        autocast_ctx = nullcontext()

    # Build per-case evaluation
    if args.validation_images is None or args.validation_images == ['None']:
        raise ValueError("validation_images must be provided and non-empty")
    if args.validation_prompt is None:
        raise ValueError("validation_prompt must be provided and non-empty")

    val_imgs = args.validation_images
    prompts = args.validation_prompt
    # Prepend instruction to each prompt (same as dataset/test requirement)
    instruction = "Fill in the white region naturally and adapt the foreground into the background. Fix the perspective of the foreground object if necessary."
    try:
        prompts = [f"{instruction} {p}".strip() if isinstance(p, str) and len(p.strip()) > 0 else instruction for p in prompts]
    except Exception:
        # Fallback: keep original prompts if unexpected
        pass

    if not (len(val_imgs) == len(prompts)):
        raise ValueError(
            f"Length mismatch: validation_images={len(val_imgs)}, validation_prompt={len(prompts)}"
        )

    results = []

    def _resize_to_preferred(img: Image.Image) -> Image.Image:
        w, h = img.size
        aspect_ratio = w / h if h != 0 else 1.0
        _, target_w, target_h = min(
            (abs(aspect_ratio - (pref_w / pref_h)), pref_w, pref_h)
            for (pref_h, pref_w) in PREFERRED_KONTEXT_RESOLUTIONS
        )
        return img.resize((target_w, target_h), Image.BICUBIC)

    # Distributed per-rank assignment: each process handles its own slice of cases
    num_cases = len(prompts)
    logger.info(f"Paired validation (distributed): {num_cases} cases across {accelerator.num_processes} ranks")

    # Indices assigned to this rank
    rank = accelerator.process_index
    world_size = accelerator.num_processes
    local_indices = list(range(rank, num_cases, world_size))

    local_images = []
    with autocast_ctx:
        for idx in local_indices:
            try:
                base_img = Image.open(val_imgs[idx]).convert("RGB")
                resized_img = _resize_to_preferred(base_img)
            except Exception as e:
                raise ValueError(f"Failed to load/resize validation image idx={idx}: {e}")

            case_args = dict(pipeline_args) if pipeline_args is not None else {}
            case_args.pop("height", None)
            case_args.pop("width", None)
            if resized_img is not None:
                tw, th = resized_img.size
                case_args["height"] = th
                case_args["width"] = tw

            case_args["prompt"] = prompts[idx]
            img = pipeline(image=resized_img, **case_args, generator=generator).images[0]
            local_images.append(img)

    # Gather all images per rank (pad to equal count) to main process
    fixed_size = (1024, 1024)
    max_local = int(math.ceil(num_cases / world_size)) if world_size > 0 else len(local_images)
    # Build per-rank batch tensors
    imgs_rank = []
    idx_rank = []
    has_rank = []
    for j in range(max_local):
        if j < len(local_images):
            resized = local_images[j].resize(fixed_size, Image.BICUBIC)
            img_np = np.asarray(resized).astype(np.uint8)
            imgs_rank.append(torch.from_numpy(img_np))
            idx_rank.append(local_indices[j])
            has_rank.append(1)
        else:
            imgs_rank.append(torch.from_numpy(np.zeros((fixed_size[1], fixed_size[0], 3), dtype=np.uint8)))
            idx_rank.append(-1)
            has_rank.append(0)
    imgs_rank_tensor = torch.stack([t.to(device=accelerator.device) for t in imgs_rank], dim=0)  # [max_local, H, W, C]
    idx_rank_tensor = torch.tensor(idx_rank, device=accelerator.device, dtype=torch.long)  # [max_local]
    has_rank_tensor = torch.tensor(has_rank, device=accelerator.device, dtype=torch.int)   # [max_local]

    gathered_has = accelerator.gather(has_rank_tensor)        # [world * max_local]
    gathered_idx = accelerator.gather(idx_rank_tensor)        # [world * max_local]
    gathered_imgs = accelerator.gather(imgs_rank_tensor)      # [world * max_local, H, W, C]

    if accelerator.is_main_process:
        world = int(world_size)
        slots = int(max_local)
        try:
            gathered_has = gathered_has.view(world, slots)
            gathered_idx = gathered_idx.view(world, slots)
            gathered_imgs = gathered_imgs.view(world, slots, fixed_size[1], fixed_size[0], 3)
        except Exception:
            # Fallback: treat as flat if reshape fails
            gathered_has = gathered_has.view(-1, 1)
            gathered_idx = gathered_idx.view(-1, 1)
            gathered_imgs = gathered_imgs.view(-1, 1, fixed_size[1], fixed_size[0], 3)
            world = int(gathered_has.shape[0])
            slots = 1
        for i in range(world):
            for j in range(slots):
                if int(gathered_has[i, j].item()) == 1:
                    idx = int(gathered_idx[i, j].item())
                    arr = gathered_imgs[i, j].cpu().numpy()
                    pil_img = Image.fromarray(arr.astype(np.uint8))
                    # Resize back to original validation image size
                    try:
                        orig = Image.open(val_imgs[idx]).convert("RGB")
                        pil_img = pil_img.resize(orig.size, Image.BICUBIC)
                    except Exception:
                        pass
                    results.append(pil_img)

    # Log results (resize to 1024x1024 for saving or external trackers). Skip TensorBoard per request.
    resized_for_log = [img.resize((1024, 1024), Image.BICUBIC) for img in results]
    for tracker in accelerator.trackers:
        phase_name = "test" if is_final_validation else "validation"
        if tracker.name == "tensorboard":
            continue
        if tracker.name == "wandb":
            tracker.log({
                phase_name: [wandb.Image(image, caption=f"{i}: {prompts[i] if i < len(prompts) else ''}") for i, image in enumerate(resized_for_log)]
            })

    del pipeline
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return results


def save_with_retry(img: Image.Image, path: str, max_retries: int = 3) -> bool:
    """Save PIL image with simple retry and exponential backoff to mitigate transient I/O errors."""
    last_err = None
    for attempt in range(max_retries):
        try:
            os.makedirs(os.path.dirname(path), exist_ok=True)
            img.save(path)
            return True
        except OSError as e:
            last_err = e
            # Exponential backoff: 1.0, 1.5, 2.25 seconds ...
            time.sleep(1.5 ** attempt)
    logger.warning(f"Failed to save {path} after {max_retries} retries: {last_err}")
    return False


def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str, revision: str, subfolder: str = "text_encoder"):
    text_encoder_config = transformers.PretrainedConfig.from_pretrained(
        pretrained_model_name_or_path, subfolder=subfolder, revision=revision
    )
    model_class = text_encoder_config.architectures[0]
    if model_class == "CLIPTextModel":
        from transformers import CLIPTextModel

        return CLIPTextModel
    elif model_class == "T5EncoderModel":
        from transformers import T5EncoderModel

        return T5EncoderModel
    else:
        raise ValueError(f"{model_class} is not supported.")


def parse_args(input_args=None):
    parser = argparse.ArgumentParser(description="Training script for Flux Kontext with EasyControl.")
    parser.add_argument("--mode", type=str, default=None, help="Controller mode; kept for compatibility.")

    # Dataset arguments
    parser.add_argument("--dataset_mode", type=str, default="mixed", choices=["interactive", "placement", "pexels", "mixed"], 
                        help="Dataset mode: interactive, placement, pexels, or mixed")
    parser.add_argument("--train_data_jsonl", type=str, default="/robby/share/Editing/lzc/HOI_v1/final_metadata.jsonl", 
                        help="Path to interactive dataset JSONL")
    parser.add_argument("--placement_data_jsonl", type=str, default="/robby/share/Editing/lzc/subject_placement/metadata_relight.jsonl", 
                        help="Path to placement dataset JSONL")
    parser.add_argument("--pexels_data_jsonl", type=str, default=None, 
                        help="Path to pexels dataset JSONL")
    parser.add_argument("--interactive_base_dir", type=str, default="/robby/share/Editing/lzc/HOI_v1", 
                        help="Base directory for interactive dataset")
    parser.add_argument("--placement_base_dir", type=str, default="/robby/share/Editing/lzc/subject_placement", 
                        help="Base directory for placement dataset")
    parser.add_argument("--pexels_base_dir", type=str, default=None, 
                        help="Base directory for pexels dataset")
    parser.add_argument("--pexels_relight_base_dir", type=str, default=None, 
                        help="Base directory for pexels relighted images")
    parser.add_argument("--seg_base_dir", type=str, default=None, 
                        help="Directory containing segmentation maps for pexels dataset")
    parser.add_argument("--interactive_weight", type=float, default=1.0, 
                        help="Sampling weight for interactive dataset (default: 1.0)")
    parser.add_argument("--placement_weight", type=float, default=1.0, 
                        help="Sampling weight for placement dataset (default: 1.0)")
    parser.add_argument("--pexels_weight", type=float, default=0.1, 
                        help="Sampling weight for pexels dataset (default: 1.0)")
    parser.add_argument("--pretrained_model_name_or_path", type=str, default="", required=False, help="Base model path")
    parser.add_argument("--pretrained_lora_path", type=str, default=None, required=False, help="LoRA checkpoint to initialize from")
    parser.add_argument("--revision", type=str, default=None, required=False, help="Revision of pretrained model")
    parser.add_argument("--variant", type=str, default=None, help="Variant of the model files")

    parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.")
    parser.add_argument("--max_sequence_length", type=int, default=128, help="Max sequence length for T5")
    parser.add_argument("--kontext", type=str, default="enable")
    parser.add_argument("--validation_prompt", type=str, nargs="+", default=None)
    parser.add_argument("--validation_images", type=str, nargs="+", default=None, help="List of valiadation images")
    parser.add_argument("--num_validation_images", type=int, default=4)
    parser.add_argument("--validation_steps", type=int, default=20)

    parser.add_argument("--ranks", type=int, nargs="+", default=[32], help="LoRA ranks")
    parser.add_argument("--output_dir", type=str, default="", help="Output directory")
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--train_batch_size", type=int, default=1)
    parser.add_argument("--num_train_epochs", type=int, default=50)
    parser.add_argument("--max_train_steps", type=int, default=None)
    parser.add_argument("--checkpointing_steps", type=int, default=1000)
    parser.add_argument("--checkpoints_total_limit", type=int, default=None)
    parser.add_argument("--resume_from_checkpoint", type=str, default=None)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--gradient_checkpointing", action="store_true")
    parser.add_argument("--learning_rate", type=float, default=1e-4)
    parser.add_argument("--guidance_scale", type=float, default=1.0, help="Flux Kontext is guidance distilled")
    parser.add_argument("--scale_lr", action="store_true", default=False)
    parser.add_argument("--lr_scheduler", type=str, default="constant")
    parser.add_argument("--lr_warmup_steps", type=int, default=500)
    parser.add_argument("--lr_num_cycles", type=int, default=1)
    parser.add_argument("--lr_power", type=float, default=1.0)
    parser.add_argument("--dataloader_num_workers", type=int, default=8)
    parser.add_argument("--weighting_scheme", type=str, default="none", choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"])
    parser.add_argument("--logit_mean", type=float, default=0.0)
    parser.add_argument("--logit_std", type=float, default=1.0)
    parser.add_argument("--mode_scale", type=float, default=1.29)
    parser.add_argument("--optimizer", type=str, default="AdamW")
    parser.add_argument("--use_8bit_adam", action="store_true")
    parser.add_argument("--adam_beta1", type=float, default=0.9)
    parser.add_argument("--adam_beta2", type=float, default=0.999)
    parser.add_argument("--prodigy_beta3", type=float, default=None)
    parser.add_argument("--prodigy_decouple", type=bool, default=True)
    parser.add_argument("--adam_weight_decay", type=float, default=1e-04)
    parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=1e-03)
    parser.add_argument("--adam_epsilon", type=float, default=1e-08)
    parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True)
    parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True)
    parser.add_argument("--max_grad_norm", type=float, default=1.0)
    parser.add_argument("--logging_dir", type=str, default="logs")
    parser.add_argument("--cache_latents", action="store_true", default=False)
    parser.add_argument("--report_to", type=str, default="tensorboard")
    parser.add_argument("--mixed_precision", type=str, default="bf16", choices=["no", "fp16", "bf16"])
    parser.add_argument("--upcast_before_saving", action="store_true", default=False)

    # Blending options for dataset pixel_values
    parser.add_argument("--blend_pixel_values", action="store_true", help="Blend target/source into pixel_values using mask")
    parser.add_argument("--blend_kernel", type=int, default=21, help="Gaussian blur kernel size (must be odd)")
    parser.add_argument("--blend_sigma", type=float, default=10.0, help="Gaussian blur sigma")

    if input_args is not None:
        args = parser.parse_args(input_args)
    else:
        args = parser.parse_args()
    return args


def main(args):
    if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
        raise ValueError("Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 or fp32 instead.")

    if args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)
        os.makedirs(args.logging_dir, exist_ok=True)
    logging_dir = Path(args.output_dir, args.logging_dir)

    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
    kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=args.report_to,
        project_config=accelerator_project_config,
        kwargs_handlers=[kwargs],
    )

    if torch.backends.mps.is_available():
        accelerator.native_amp = False

    if args.report_to == "wandb":
        if not is_wandb_available():
            raise ImportError("Install wandb for logging during training.")

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        transformers.utils.logging.set_verbosity_warning()
        diffusers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    if args.seed is not None:
        set_seed(args.seed)

    if accelerator.is_main_process and args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)

    # Tokenizers
    tokenizer_one = transformers.CLIPTokenizer.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
    )
    tokenizer_two = transformers.T5TokenizerFast.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="tokenizer_2", revision=args.revision
    )

    # Text encoders
    text_encoder_cls_one = import_model_class_from_model_name_or_path(args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder")
    text_encoder_cls_two = import_model_class_from_model_name_or_path(args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder_2")

    # Scheduler and models
    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
    noise_scheduler_copy = copy.deepcopy(noise_scheduler)
    text_encoder_one, text_encoder_two = load_text_encoders(args, text_encoder_cls_one, text_encoder_cls_two)
    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, variant=args.variant)
    transformer = FluxTransformer2DModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant)

    # Train only LoRA adapters: freeze base transformer/text encoders/vae
    transformer.requires_grad_(False)
    vae.requires_grad_(False)
    text_encoder_one.requires_grad_(False)
    text_encoder_two.requires_grad_(False)

    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16

    if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
        raise ValueError("Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 or fp32 instead.")

    vae.to(accelerator.device, dtype=weight_dtype)
    transformer.to(accelerator.device, dtype=weight_dtype)
    text_encoder_one.to(accelerator.device, dtype=weight_dtype)
    text_encoder_two.to(accelerator.device, dtype=weight_dtype)

    if args.gradient_checkpointing:
        transformer.enable_gradient_checkpointing()

    # Setup standard PEFT LoRA on FluxTransformer2DModel
    # target_modules = [
    #     "attn.to_k",
    #     "attn.to_q",
    #     "attn.to_v",
    #     "attn.to_out.0",
    #     "attn.add_k_proj",
    #     "attn.add_q_proj",
    #     "attn.add_v_proj",
    #     "attn.to_add_out",
    #     "ff.net.0.proj",
    #     "ff.net.2",
    #     "ff_context.net.0.proj",
    #     "ff_context.net.2",
    # ]
    target_modules = [
        "attn.to_k",
        "attn.to_q",
        "attn.to_v",
        "attn.to_out.0",
        "attn.add_k_proj",
        "attn.add_q_proj",
        "attn.add_v_proj",
        "attn.to_add_out",
        "ff.net.0.proj",
        "ff.net.2",
        "ff_context.net.0.proj",
        "ff_context.net.2",
        # ===========================================================
        # 【补全部分 1】: 单流模块 (single_transformer_blocks) 的专属层
        # ===========================================================
        # 说明：单流块中的注意力层 (to_q, to_k, to_v) 已被上面的通用名称覆盖。
        # 这里补充的是它们特有的 MLP 和输出层。
        "proj_mlp",
        "proj_out", # 这个名称也会匹配单流块各自的输出层和模型总输出层

        # ===========================================================
        # 【补全部分 2】: 所有的归一化 (Norm) 层
        # ===========================================================
        # 说明：这些层负责调整特征分布，对风格学习很重要。
        # 使用 "linear" 可以一次性匹配所有以 ".linear" 结尾的Norm层。
        "linear", # 匹配 norm1.linear, norm1_context.linear, norm.linear, norm_out.linear
    ]
    lora_rank = int(args.ranks[0]) if isinstance(args.ranks, list) and len(args.ranks) > 0 else 256
    lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=lora_rank,
        init_lora_weights="gaussian",
        target_modules=target_modules,
    )
    transformer.add_adapter(lora_config)
    transformer.train()
    print(sum([p.numel() for p in transformer.parameters() if p.requires_grad]) / 1000000, 'M parameters')

    def unwrap_model(model):
        model = accelerator.unwrap_model(model)
        model = model._orig_mod if is_compiled_module(model) else model
        return model

    if args.resume_from_checkpoint:
        path = args.resume_from_checkpoint
        global_step = int(path.split("-")[-1])
        initial_global_step = global_step
    else:
        initial_global_step = 0
        global_step = 0
        first_epoch = 0

    if args.scale_lr:
        args.learning_rate = (
            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
        )

    if args.mixed_precision == "fp16":
        models = [transformer]
        cast_training_params(models, dtype=torch.float32)

    params_to_optimize = [p for p in transformer.parameters() if p.requires_grad]
    transformer_parameters_with_lr = {"params": params_to_optimize, "lr": args.learning_rate}
    # print(sum([p.numel() for p in transformer.parameters() if p.requires_grad]) / 1000000, 'parameters')

    optimizer_class = torch.optim.AdamW
    optimizer = optimizer_class(
        [transformer_parameters_with_lr],
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    tokenizers = [tokenizer_one, tokenizer_two]
    text_encoders = [text_encoder_one, text_encoder_two]

    # Create dataset based on mode
    if args.dataset_mode == "mixed":
        # Mixed mode: combine all available datasets
        train_dataset = make_mixed_dataset(
            args,
            tokenizers,
            interactive_jsonl_path=args.train_data_jsonl,
            placement_jsonl_path=args.placement_data_jsonl,
            pexels_jsonl_path=args.pexels_data_jsonl,
            interactive_base_dir=args.interactive_base_dir,
            placement_base_dir=args.placement_base_dir,
            pexels_base_dir=args.pexels_base_dir,
            interactive_weight=args.interactive_weight,
            placement_weight=args.placement_weight,
            pexels_weight=args.pexels_weight,
            accelerator=accelerator
        )
        weights_str = []
        if args.train_data_jsonl:
            weights_str.append(f"Interactive: {args.interactive_weight:.2f}")
        if args.placement_data_jsonl:
            weights_str.append(f"Placement: {args.placement_weight:.2f}")
        if args.pexels_data_jsonl:
            weights_str.append(f"Pexels: {args.pexels_weight:.2f}")
        logger.info(f"Mixed dataset created with weights - {', '.join(weights_str)}")
    elif args.dataset_mode == "pexels":
        if not args.pexels_data_jsonl:
            raise ValueError("pexels_data_jsonl must be provided for pexels mode")
        train_dataset = make_pexels_dataset_subjects(args, tokenizers, accelerator)
    elif args.dataset_mode == "placement":
        if not args.placement_data_jsonl:
            raise ValueError("placement_data_jsonl must be provided for placement mode")
        train_dataset = make_placement_dataset_subjects(args, tokenizers, accelerator)
    else:  # interactive mode
        train_dataset = make_interactive_dataset_subjects(args, tokenizers, accelerator)
    
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=args.train_batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=args.dataloader_num_workers,
    )

    vae_config_shift_factor = vae.config.shift_factor
    vae_config_scaling_factor = vae.config.scaling_factor

    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if args.resume_from_checkpoint:
        first_epoch = global_step // num_update_steps_per_epoch
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True

    lr_scheduler = get_scheduler(
        args.lr_scheduler,
        optimizer=optimizer,
        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
        num_training_steps=args.max_train_steps * accelerator.num_processes,
        num_cycles=args.lr_num_cycles,
        power=args.lr_power,
    )

    transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        transformer, optimizer, train_dataloader, lr_scheduler
    )

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if overrode_max_train_steps:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    # Sanitize config for TensorBoard hparams (only allow int/float/bool/str/tensor). Others are stringified if possible; otherwise dropped
    def _sanitize_hparams(config_dict):
        sanitized = {}
        for key, value in dict(config_dict).items():
            try:
                if value is None:
                    continue
                # numpy scalar types
                if isinstance(value, (np.integer,)):
                    sanitized[key] = int(value)
                elif isinstance(value, (np.floating,)):
                    sanitized[key] = float(value)
                elif isinstance(value, (int, float, bool, str)):
                    sanitized[key] = value
                elif isinstance(value, Path):
                    sanitized[key] = str(value)
                elif isinstance(value, (list, tuple)):
                    # stringify simple sequences; skip if fails
                    sanitized[key] = str(value)
                else:
                    # best-effort stringify
                    sanitized[key] = str(value)
            except Exception:
                # skip unconvertible entries
                continue
        return sanitized

    if accelerator.is_main_process:
        tracker_name = "Easy_Control_Kontext"
        accelerator.init_trackers(tracker_name, config=_sanitize_hparams(vars(args)))

    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")

    progress_bar = tqdm(
        range(0, args.max_train_steps),
        initial=initial_global_step,
        desc="Steps",
        disable=not accelerator.is_local_main_process,
    )

    def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
        sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
        schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device)
        timesteps = timesteps.to(accelerator.device)
        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
        sigma = sigmas[step_indices].flatten()
        while len(sigma.shape) < n_dim:
            sigma = sigma.unsqueeze(-1)
        return sigma

    # Kontext specifics
    vae_scale_factor = 8  # Kontext uses 8x VAE factor; pack/unpack uses additional 2x in methods

    for epoch in range(first_epoch, args.num_train_epochs):
        transformer.train()
        for step, batch in enumerate(train_dataloader):
            models_to_accumulate = [transformer]
            with accelerator.accumulate(models_to_accumulate):
                tokens = [batch["text_ids_1"], batch["text_ids_2"]]
                prompt_embeds, pooled_prompt_embeds, text_ids = encode_token_ids(text_encoders, tokens, accelerator)
                prompt_embeds = prompt_embeds.to(dtype=vae.dtype, device=accelerator.device)
                pooled_prompt_embeds = pooled_prompt_embeds.to(dtype=vae.dtype, device=accelerator.device)
                text_ids = text_ids.to(dtype=vae.dtype, device=accelerator.device)

                pixel_values = batch["pixel_values"].to(dtype=vae.dtype)
                height_ = 2 * (int(pixel_values.shape[-2]) // (vae_scale_factor * 2))
                width_ = 2 * (int(pixel_values.shape[-1]) // (vae_scale_factor * 2))

                model_input = vae.encode(pixel_values).latent_dist.sample()
                model_input = (model_input - vae_config_shift_factor) * vae_config_scaling_factor
                model_input = model_input.to(dtype=weight_dtype)

                # Prepare latent ids for transformer (positional encodings)
                latent_image_ids = FluxKontextPipeline._prepare_latent_image_ids(
                    batch_size=model_input.shape[0],
                    height=model_input.shape[2] // 2,
                    width=model_input.shape[3] // 2,
                    device=accelerator.device,
                    dtype=weight_dtype,
                )

                noise = torch.randn_like(model_input)
                bsz = model_input.shape[0]

                u = compute_density_for_timestep_sampling(
                    weighting_scheme=args.weighting_scheme,
                    batch_size=bsz,
                    logit_mean=args.logit_mean,
                    logit_std=args.logit_std,
                    mode_scale=args.mode_scale,
                )
                indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
                timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device)

                sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
                noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise

                packed_noisy_model_input = FluxKontextPipeline._pack_latents(
                    noisy_model_input,
                    batch_size=model_input.shape[0],
                    num_channels_latents=model_input.shape[1],
                    height=model_input.shape[2],
                    width=model_input.shape[3],
                )
                
                if accelerator.unwrap_model(transformer).config.guidance_embeds:
                    guidance = torch.tensor([args.guidance_scale], device=accelerator.device)
                    guidance = guidance.expand(model_input.shape[0])
                else:
                    guidance = None
                
                # If kontext editing is enabled, append source image latents to the sequence
                latent_model_input = packed_noisy_model_input
                if args.kontext == "enable":
                    source_pixel_values = batch["source_pixel_values"].to(dtype=vae.dtype)
                    source_image_latents = vae.encode(source_pixel_values).latent_dist.sample()
                    source_image_latents = (source_image_latents - vae_config_shift_factor) * vae_config_scaling_factor
                    image_latent_h, image_latent_w = source_image_latents.shape[2:]
                    packed_image_latents = FluxKontextPipeline._pack_latents(
                        source_image_latents,
                        batch_size=source_image_latents.shape[0],
                        num_channels_latents=source_image_latents.shape[1],
                        height=image_latent_h,
                        width=image_latent_w,
                    )
                    source_image_ids = FluxKontextPipeline._prepare_latent_image_ids(
                        batch_size=source_image_latents.shape[0],
                        height=image_latent_h // 2,
                        width=image_latent_w // 2,
                        device=accelerator.device,
                        dtype=weight_dtype,
                    )
                    source_image_ids[..., 0] = 1
                    latent_model_input = torch.cat([latent_model_input, packed_image_latents], dim=1)
                    latent_image_ids = torch.cat([latent_image_ids, source_image_ids], dim=0)

                # Forward transformer with packed latents and ids
                model_pred = transformer(
                    hidden_states=latent_model_input,
                    timestep=timesteps / 1000,
                    guidance=guidance,
                    pooled_projections=pooled_prompt_embeds,
                    encoder_hidden_states=prompt_embeds,
                    txt_ids=text_ids,
                    img_ids=latent_image_ids,
                    return_dict=False,
                )[0]

                model_pred = model_pred[:, : packed_noisy_model_input.size(1)]

                model_pred = FluxKontextPipeline._unpack_latents(
                    model_pred,
                    height=int(pixel_values.shape[-2]),
                    width=int(pixel_values.shape[-1]),
                    vae_scale_factor=vae_scale_factor,
                )

                weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)
                target = noise - model_input

                loss = torch.mean((weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1), 1)
                loss = loss.mean()
                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    params_to_clip = (transformer.parameters())
                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)

                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.sync_gradients:
                progress_bar.update(1)
                global_step += 1

                if accelerator.is_main_process:
                    if global_step % args.checkpointing_steps == 0:
                        if args.checkpoints_total_limit is not None:
                            checkpoints = os.listdir(args.output_dir)
                            checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
                            checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
                            if len(checkpoints) >= args.checkpoints_total_limit:
                                num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
                                removing_checkpoints = checkpoints[0:num_to_remove]
                                logger.info(f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints")
                                logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
                                for removing_checkpoint in removing_checkpoints:
                                    removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
                                    shutil.rmtree(removing_checkpoint)

                        save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                        os.makedirs(save_path, exist_ok=True)
                        unwrapped = accelerator.unwrap_model(transformer)
                        peft_state = get_peft_model_state_dict(unwrapped)
                        # Convert PEFT state dict to diffusers LoRA format for transformer
                        diffusers_lora = convert_state_dict_to_diffusers(peft_state)
                        save_file(diffusers_lora, os.path.join(save_path, "pytorch_lora_weights.safetensors"))
                        logger.info(f"Saved state to {save_path}")

            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
            progress_bar.set_postfix(**logs)
            accelerator.log(logs, step=global_step)

            if args.validation_prompt is not None and global_step % args.validation_steps == 0:
                # Create pipeline on every rank to run validation in parallel
                pipeline = FluxKontextPipeline.from_pretrained(
                    args.pretrained_model_name_or_path,
                    vae=vae,
                    text_encoder=accelerator.unwrap_model(text_encoder_one),
                    text_encoder_2=accelerator.unwrap_model(text_encoder_two),
                    transformer=accelerator.unwrap_model(transformer),
                    revision=args.revision,
                    variant=args.variant,
                    torch_dtype=weight_dtype,
                )

                pipeline_args = {
                    "prompt": args.validation_prompt,
                    "guidance_scale": 3.5,
                    "num_inference_steps": 20,
                    "max_sequence_length": 128,
                }

                images = log_validation(
                    pipeline=pipeline,
                    args=args,
                    accelerator=accelerator,
                    pipeline_args=pipeline_args,
                    step=global_step,
                    torch_dtype=weight_dtype,
                )

                # Only main process saves/logs
                if accelerator.is_main_process:
                    save_path = os.path.join(args.output_dir, "validation")
                    os.makedirs(save_path, exist_ok=True)
                    save_folder = os.path.join(save_path, f"checkpoint-{global_step}")
                    os.makedirs(save_folder, exist_ok=True)
                    for idx, img in enumerate(images):
                        out_path = os.path.join(save_folder, f"{idx}.jpg")
                        save_with_retry(img, out_path)
                del pipeline

    accelerator.wait_for_everyone()
    accelerator.end_training()


if __name__ == "__main__":
    args = parse_args()
    main(args)