train/train_kontext_local.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_local import make_train_dataset_mixed, 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 compute_background_preserving_loss(model_pred, target, mask_values, weighting, background_weight: float = 3.0):
    """
    Compute loss with higher penalty on background (non-masked) regions to preserve them.
    model_pred/target: [B, C, H, W]
    mask_values: [B, 1, H_img, W_img] with values in {0,1} at image resolution
    weighting: broadcastable to [B, C, H, W]
    Returns per-pixel loss map [B, C, H, W]
    """
    base_loss = (weighting.float() * (model_pred.float() - target.float()) ** 2)
    mask_latent = torch.nn.functional.interpolate(
        mask_values,
        size=(model_pred.shape[2], model_pred.shape[3]),
        mode='bilinear',
        align_corners=False,
    )
    foreground_mask = mask_latent
    background_mask = 1.0 - mask_latent
    foreground_mask = foreground_mask.expand_as(base_loss)
    background_mask = background_mask.expand_as(base_loss)
    foreground_loss = base_loss * foreground_mask
    background_loss = base_loss * background_mask * float(background_weight)
    total_loss = foreground_loss + background_loss
    return total_loss

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)

    # 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")

    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_control = dict(case_args.get("control_dict", {}))
            spatial_case = spatial_ls[idx]

            # Compose masked image cond: resized_img * (1 - binary_mask)
            try:
                mask_img = Image.open(spatial_case).convert("L") if isinstance(spatial_case, str) else spatial_case.convert("L")
            except Exception:
                mask_img = spatial_case.convert("L")
            mask_img = mask_img.resize(resized_img.size, Image.NEAREST)
            mask_np = np.array(mask_img)
            mask_bin = (mask_np > 127).astype(np.uint8)
            inv_mask = (1 - mask_bin).astype(np.uint8)
            base_np = np.array(resized_img)
            masked_np = base_np * inv_mask[..., None]
            masked_img = Image.fromarray(masked_np.astype(np.uint8))

            case_control["spatial_images"] = [masked_img]
            case_args["control_dict"] = case_control

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

    # Gather one image per rank (pad missing ranks with black images) to main process
    fixed_size = (1024, 1024)
    has_sample = torch.tensor([1 if len(local_images) > 0 else 0], device=accelerator.device, dtype=torch.int)
    local_idx = torch.tensor([local_indices[0] if len(local_indices) > 0 else -1], device=accelerator.device, dtype=torch.long)
    if len(local_images) > 0:
        gathered_img = local_images[0].resize(fixed_size, Image.BICUBIC)
        img_np = np.asarray(gathered_img).astype(np.uint8)
    else:
        img_np = np.zeros((fixed_size[1], fixed_size[0], 3), dtype=np.uint8)
    img_tensor = torch.from_numpy(img_np).to(device=accelerator.device)
    if img_tensor.ndim == 3:
        img_tensor = img_tensor.unsqueeze(0)

    gathered_has = accelerator.gather(has_sample)
    gathered_idx = accelerator.gather(local_idx)
    gathered_imgs = accelerator.gather(img_tensor)

    if accelerator.is_main_process:
        for i in range(int(gathered_has.shape[0])):
            if int(gathered_has[i].item()) == 1:
                idx = int(gathered_idx[i].item())
                arr = gathered_imgs[i].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)

    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.")

    # New dataset (local edits + inpaint JSONL) mixed 1:1
    parser.add_argument("--local_edits_json", type=str, default="/robby/share/Editing/qingyan/InstructV2V/Qwen2_5_72B_instructs_10W.json", help="Path to local edits JSON")
    parser.add_argument("--train_data_dir", type=str, default="/robby/share/Editing/lzc/data/pexel_final/inpaint_edit_outputs_merged.jsonl", help="Path to inpaint JSONL file for mixing 1:1")
    parser.add_argument("--source_frames_dir", type=str, default="/robby/share/Editing/qingyan/InstructV2V/pexel-video-merged-1frame", help="Root dir containing group folders like 0139")
    parser.add_argument("--target_frames_dir", type=str, default="/robby/share/Editing/qingyan/InstructV2V/pexel-video-1frame-kontext-edit/local", help="Root dir containing group folders like 0139")
    parser.add_argument("--masks_dir", type=str, default="/robby/share/Editing/lzc/InstructV2V/diff_masks", help="Root dir of precomputed masks organized as <group>/<prefix>_{i}.png")
    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("--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=[256], help="LoRA ranks")
    parser.add_argument("--network_alphas", type=int, nargs="+", default=[256], 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=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)
    parser.add_argument("--mix_ratio", type=float, default=0, help="Ratio of inpaint to local edits (B per A). 0=only local edits, 1=1:1, 2=1:2")
    parser.add_argument("--background_weight", type=float, default=1.0, help="Background preserving loss weight multiplier")

    # 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
    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_mixed(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

                # mask_values = batch.get("mask_values")
                # if mask_values is not None:
                #     mask_values = mask_values.to(device=accelerator.device, dtype=model_pred.dtype)
                #     loss_map = compute_background_preserving_loss(
                #         model_pred=model_pred,
                #         target=target,
                #         mask_values=mask_values,
                #         weighting=weighting,
                #         background_weight=args.background_weight,
                #     )
                #     loss = torch.mean(loss_map.reshape(target.shape[0], -1), 1)
                #     loss = loss.mean()
                # else:
                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 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.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},
                }

                images = log_validation(
                    pipeline=pipeline,
                    args=args,
                    accelerator=accelerator,
                    pipeline_args=pipeline_args,
                    step=global_step,
                    torch_dtype=weight_dtype,
                )
                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):
                        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)