train/train_kontext_edge.py

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

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
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_edge import make_train_dataset_inpaint_mask, collate_fn
from src.pipeline_flux_kontext_control import (
    FluxKontextControlPipeline,
    resize_position_encoding,
    prepare_latent_subject_ids,
    PREFERRED_KONTEXT_RESOLUTIONS
)
from src.transformer_flux import FluxTransformer2DModel
from diffusers.models.attention_processor import FluxAttnProcessor2_0
from src.layers import MultiDoubleStreamBlockLoraProcessor, MultiSingleStreamBlockLoraProcessor
from tqdm.auto import tqdm

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__)


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

    generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
    autocast_ctx = nullcontext()

    # Build per-case evaluation: require equal lengths for image, spatial image, and prompt
    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")

    control_dict_root = dict(pipeline_args.get("control_dict", {})) if pipeline_args is not None else {}
    spatial_ls = control_dict_root.get("spatial_images", []) or []

    val_imgs = args.validation_images
    prompts = args.validation_prompt

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

    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)

    # Strict per-case loop
    num_cases = len(prompts)
    logger.info(f"Paired validation: {num_cases} (image, spatial, prompt) cases")
    with autocast_ctx:
        for idx in range(num_cases):
            resized_img = None
            # If validation image path is a non-empty string, load and resize; otherwise, skip passing image
            if isinstance(val_imgs[idx], str) and val_imgs[idx] != "":
                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
            else:
                # When no image is provided, default to 1024x1024
                case_args["height"] = 1024
                case_args["width"] = 1024

            # Bind single spatial control image per case; pass it directly (no masking)
            case_control = dict(case_args.get("control_dict", {}))
            spatial_case = spatial_ls[idx]

            # Load spatial image if it's a path; else assume it's already an image
            try:
                spatial_img = Image.open(spatial_case).convert("RGB") if isinstance(spatial_case, str) else spatial_case
            except Exception:
                spatial_img = spatial_case

            case_control["spatial_images"] = [spatial_img]
            case_control["subject_images"] = []
            case_args["control_dict"] = case_control

            # Override prompt per case
            case_args["prompt"] = prompts[idx]

            if resized_img is not None:
                img = pipeline(image=resized_img, **case_args, generator=generator).images[0]
            else:
                img = pipeline(**case_args, generator=generator).images[0]
            results.append(img)

    # Log results (resize to 1024x1024 for logging only)
    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":
            np_images = np.stack([np.asarray(img) for img in resized_for_log])
            tracker.writer.add_images(phase_name, np_images, step, dataformats="NHWC")
        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 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("--lora_num", type=int, default=1, help="number of the lora.")
    parser.add_argument("--cond_size", type=int, default=512, help="size of the condition data.")
    parser.add_argument("--mode", type=str, default=None, help="Controller mode; kept for compatibility.")

    parser.add_argument("--train_data_dir", type=str, default="", help="Path to JSONL dataset.")
    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="disable")
    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("--subject_test_images", type=str, nargs="+", default=None, help="List of subject test images")
    parser.add_argument("--spatial_test_images", type=str, nargs="+", default=None, help="List of spatial test 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=[128], help="LoRA ranks")
    parser.add_argument("--network_alphas", type=int, nargs="+", default=[128], help="LoRA network alphas")
    parser.add_argument("--output_dir", type=str, default="/tiamat-NAS/zhangyuxuan/projects2/Easy_Control_0120/single_models/subject_model", 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=1)
    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)

    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
    transformer.requires_grad_(True)
    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 LoRA attention processors
    if args.pretrained_lora_path is not None:
        lora_path = args.pretrained_lora_path
        checkpoint = load_checkpoint(lora_path)
        lora_attn_procs = {}
        double_blocks_idx = list(range(19))
        single_blocks_idx = list(range(38))
        number = 1
        for name, attn_processor in transformer.attn_processors.items():
            match = re.search(r'\.(\d+)\.', name)
            if match:
                layer_index = int(match.group(1))
            if name.startswith("transformer_blocks") and layer_index in double_blocks_idx:
                lora_state_dicts = {}
                for key, value in checkpoint.items():
                    if re.search(r'\.(\d+)\.', key):
                        checkpoint_layer_index = int(re.search(r'\.(\d+)\.', key).group(1))
                        if checkpoint_layer_index == layer_index and key.startswith("transformer_blocks"):
                            lora_state_dicts[key] = value
                lora_attn_procs[name] = MultiDoubleStreamBlockLoraProcessor(
                    dim=3072, ranks=args.ranks, network_alphas=args.network_alphas, lora_weights=[1 for _ in range(args.lora_num)], device=accelerator.device, dtype=weight_dtype, cond_width=args.cond_size, cond_height=args.cond_size, n_loras=args.lora_num
                )
                for n in range(number):
                    lora_attn_procs[name].q_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.down.weight', None)
                    lora_attn_procs[name].q_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.up.weight', None)
                    lora_attn_procs[name].k_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.down.weight', None)
                    lora_attn_procs[name].k_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.up.weight', None)
                    lora_attn_procs[name].v_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.down.weight', None)
                    lora_attn_procs[name].v_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.up.weight', None)
                    lora_attn_procs[name].proj_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.proj_loras.{n}.down.weight', None)
                    lora_attn_procs[name].proj_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.proj_loras.{n}.up.weight', None)
            elif name.startswith("single_transformer_blocks") and layer_index in single_blocks_idx:
                lora_state_dicts = {}
                for key, value in checkpoint.items():
                    if re.search(r'\.(\d+)\.', key):
                        checkpoint_layer_index = int(re.search(r'\.(\d+)\.', key).group(1))
                        if checkpoint_layer_index == layer_index and key.startswith("single_transformer_blocks"):
                            lora_state_dicts[key] = value
                lora_attn_procs[name] = MultiSingleStreamBlockLoraProcessor(
                    dim=3072, ranks=args.ranks, network_alphas=args.network_alphas, lora_weights=[1 for _ in range(args.lora_num)], device=accelerator.device, dtype=weight_dtype, cond_width=args.cond_size, cond_height=args.cond_size, n_loras=args.lora_num
                )
                for n in range(number):
                    lora_attn_procs[name].q_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.down.weight', None)
                    lora_attn_procs[name].q_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.q_loras.{n}.up.weight', None)
                    lora_attn_procs[name].k_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.down.weight', None)
                    lora_attn_procs[name].k_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.k_loras.{n}.up.weight', None)
                    lora_attn_procs[name].v_loras[n].down.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.down.weight', None)
                    lora_attn_procs[name].v_loras[n].up.weight.data = lora_state_dicts.get(f'{name}.v_loras.{n}.up.weight', None)
            else:
                lora_attn_procs[name] = FluxAttnProcessor2_0()
    else:
        lora_attn_procs = {}
        double_blocks_idx = list(range(19))
        single_blocks_idx = list(range(38))
        for name, attn_processor in transformer.attn_processors.items():
            match = re.search(r'\.(\d+)\.', name)
            if match:
                layer_index = int(match.group(1))
            if name.startswith("transformer_blocks") and layer_index in double_blocks_idx:
                lora_attn_procs[name] = MultiDoubleStreamBlockLoraProcessor(
                    dim=3072, ranks=args.ranks, network_alphas=args.network_alphas, lora_weights=[1 for _ in range(args.lora_num)], device=accelerator.device, dtype=weight_dtype, cond_width=args.cond_size, cond_height=args.cond_size, n_loras=args.lora_num
                )
            elif name.startswith("single_transformer_blocks") and layer_index in single_blocks_idx:
                lora_attn_procs[name] = MultiSingleStreamBlockLoraProcessor(
                    dim=3072, ranks=args.ranks, network_alphas=args.network_alphas, lora_weights=[1 for _ in range(args.lora_num)], device=accelerator.device, dtype=weight_dtype, cond_width=args.cond_size, cond_height=args.cond_size, n_loras=args.lora_num
                )
            else:
                lora_attn_procs[name] = attn_processor

    transformer.set_attn_processor(lora_attn_procs)
    transformer.train()
    for n, param in transformer.named_parameters():
        if '_lora' not in n:
            param.requires_grad = False
    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]

    train_dataset = make_train_dataset_inpaint_mask(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
    # Match pipeline's prepare_latents cond resolution: 2 * (cond_size // (vae_scale_factor * 2))
    height_cond = 2 * (args.cond_size // (vae_scale_factor * 2))
    width_cond = 2 * (args.cond_size // (vae_scale_factor * 2))
    offset = 64

    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)

                latent_image_ids, cond_latent_image_ids = resize_position_encoding(
                    model_input.shape[0], height_, width_, height_cond, width_cond, accelerator.device, 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 = FluxKontextControlPipeline._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],
                )

                latent_image_ids_to_concat = [latent_image_ids]
                packed_cond_model_input_to_concat = []

                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 = FluxKontextControlPipeline._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 = FluxKontextControlPipeline._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  # Mark as condition
                    latent_image_ids_to_concat.append(source_image_ids)

                
                subject_pixel_values = batch.get("subject_pixel_values")
                if subject_pixel_values is not None:
                    subject_pixel_values = subject_pixel_values.to(dtype=vae.dtype)
                    subject_input = vae.encode(subject_pixel_values).latent_dist.sample()
                    subject_input = (subject_input - vae_config_shift_factor) * vae_config_scaling_factor
                    subject_input = subject_input.to(dtype=weight_dtype)
                    sub_number = subject_pixel_values.shape[-2] // args.cond_size
                    latent_subject_ids = prepare_latent_subject_ids(height_cond // 2, width_cond // 2, accelerator.device, weight_dtype)
                    latent_subject_ids[..., 0] = 2
                    latent_subject_ids[:, 1] += offset
                    sub_latent_image_ids = torch.cat([latent_subject_ids for _ in range(sub_number)], dim=0)
                    latent_image_ids_to_concat.append(sub_latent_image_ids)

                    packed_subject_model_input = FluxKontextControlPipeline._pack_latents(
                        subject_input,
                        batch_size=subject_input.shape[0],
                        num_channels_latents=subject_input.shape[1],
                        height=subject_input.shape[2],
                        width=subject_input.shape[3],
                    )
                    packed_cond_model_input_to_concat.append(packed_subject_model_input)

                cond_pixel_values = batch.get("cond_pixel_values")
                if cond_pixel_values is not None:
                    cond_pixel_values = cond_pixel_values.to(dtype=vae.dtype)
                    cond_input = vae.encode(cond_pixel_values).latent_dist.sample()
                    cond_input = (cond_input - vae_config_shift_factor) * vae_config_scaling_factor
                    cond_input = cond_input.to(dtype=weight_dtype)
                    cond_number = cond_pixel_values.shape[-2] // args.cond_size
                    cond_latent_image_ids[..., 0] = 2
                    cond_latent_image_ids_rep = torch.cat([cond_latent_image_ids for _ in range(cond_number)], dim=0)
                    latent_image_ids_to_concat.append(cond_latent_image_ids_rep)

                    packed_cond_model_input = FluxKontextControlPipeline._pack_latents(
                        cond_input,
                        batch_size=cond_input.shape[0],
                        num_channels_latents=cond_input.shape[1],
                        height=cond_input.shape[2],
                        width=cond_input.shape[3],
                    )
                    packed_cond_model_input_to_concat.append(packed_cond_model_input)

                latent_image_ids = torch.cat(latent_image_ids_to_concat, dim=0)
                cond_packed_noisy_model_input = torch.cat(packed_cond_model_input_to_concat, dim=1)

                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

                latent_model_input=packed_noisy_model_input
                if args.kontext == "enable":
                    latent_model_input = torch.cat([latent_model_input, packed_image_latents], dim=1)
                model_pred = transformer(
                    hidden_states=latent_model_input,
                    cond_hidden_states=cond_packed_noisy_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 = FluxKontextControlPipeline._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_model_state = accelerator.unwrap_model(transformer).state_dict()
                        lora_state_dict = {k: unwrapped_model_state[k] for k in unwrapped_model_state.keys() if '_lora' in k}
                        save_file(lora_state_dict, os.path.join(save_path, "lora.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 accelerator.is_main_process:
                if args.validation_prompt is not None and global_step % args.validation_steps == 0:
                    pipeline = FluxKontextControlPipeline.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,
                    )

                    if args.subject_test_images is not None and len(args.subject_test_images) != 0 and args.subject_test_images != ['None']:
                        subject_paths = args.subject_test_images
                        subject_ls = [Image.open(image_path).convert("RGB") for image_path in subject_paths]
                    else:
                        subject_ls = []
                    if args.spatial_test_images is not None and len(args.spatial_test_images) != 0 and args.spatial_test_images != ['None']:
                        spatial_paths = args.spatial_test_images
                        spatial_ls = [Image.open(image_path).convert("RGB") for image_path in spatial_paths]
                    else:
                        spatial_ls = []
                        
                    pipeline_args = {
                        "prompt": args.validation_prompt,
                        "cond_size": args.cond_size,
                        "guidance_scale": 3.5,
                        "num_inference_steps": 20,
                        "max_sequence_length": 128,
                        "control_dict": {"spatial_images": spatial_ls, "subject_images": subject_ls},
                    }

                    images = log_validation(
                        pipeline=pipeline,
                        args=args,
                        accelerator=accelerator,
                        pipeline_args=pipeline_args,
                        step=global_step,
                        torch_dtype=weight_dtype,
                    )
                    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):
                        img.save(os.path.join(save_folder, f"{idx}.jpg"))
                    del pipeline

    accelerator.wait_for_everyone()
    accelerator.end_training()


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