Edit2Perceive/models/utils.py

import torch, os
from safetensors import safe_open
from contextlib import contextmanager
import hashlib

@contextmanager
def init_weights_on_device(device = torch.device("meta"), include_buffers :bool = False):
    
    old_register_parameter = torch.nn.Module.register_parameter
    if include_buffers:
        old_register_buffer = torch.nn.Module.register_buffer
    
    def register_empty_parameter(module, name, param):
        old_register_parameter(module, name, param)
        if param is not None:
            param_cls = type(module._parameters[name])
            kwargs = module._parameters[name].__dict__
            kwargs["requires_grad"] = param.requires_grad
            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)

    def register_empty_buffer(module, name, buffer, persistent=True):
        old_register_buffer(module, name, buffer, persistent=persistent)
        if buffer is not None:
            module._buffers[name] = module._buffers[name].to(device)
            
    def patch_tensor_constructor(fn):
        def wrapper(*args, **kwargs):
            kwargs["device"] = device
            return fn(*args, **kwargs)

        return wrapper
    
    if include_buffers:
        tensor_constructors_to_patch = {
            torch_function_name: getattr(torch, torch_function_name)
            for torch_function_name in ["empty", "zeros", "ones", "full"]
        }
    else:
        tensor_constructors_to_patch = {}
    
    try:
        torch.nn.Module.register_parameter = register_empty_parameter
        if include_buffers:
            torch.nn.Module.register_buffer = register_empty_buffer
        for torch_function_name in tensor_constructors_to_patch.keys():
            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
        yield
    finally:
        torch.nn.Module.register_parameter = old_register_parameter
        if include_buffers:
            torch.nn.Module.register_buffer = old_register_buffer
        for torch_function_name, old_torch_function in tensor_constructors_to_patch.items():
            setattr(torch, torch_function_name, old_torch_function)

def load_state_dict_from_folder(file_path, torch_dtype=None):
    state_dict = {}
    for file_name in sorted(os.listdir(file_path)):
        if "." in file_name and file_name.split(".")[-1] in [
            "safetensors", "bin", "ckpt", "pth", "pt"
        ]:
            state_dict.update(load_state_dict(os.path.join(file_path, file_name), torch_dtype=torch_dtype))
    return state_dict


def load_state_dict(file_path, torch_dtype=None, device="cpu"):
    if file_path.endswith(".safetensors"):
        return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device)
    else:
        return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device)


def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"):
    state_dict = {}
    with safe_open(file_path, framework="pt", device=str(device)) as f:
        for k in f.keys():
            state_dict[k] = f.get_tensor(k)
            if torch_dtype is not None:
                state_dict[k] = state_dict[k].to(torch_dtype)
    return state_dict


def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"):
    state_dict = torch.load(file_path, map_location=device, weights_only=True)
    if torch_dtype is not None:
        for i in state_dict:
            if isinstance(state_dict[i], torch.Tensor):
                state_dict[i] = state_dict[i].to(torch_dtype)
    return state_dict


def search_for_embeddings(state_dict):
    embeddings = []
    for k in state_dict:
        if isinstance(state_dict[k], torch.Tensor):
            embeddings.append(state_dict[k])
        elif isinstance(state_dict[k], dict):
            embeddings += search_for_embeddings(state_dict[k])
    return embeddings


def search_parameter(param, state_dict):
    for name, param_ in state_dict.items():
        if param.numel() == param_.numel():
            if param.shape == param_.shape:
                if torch.dist(param, param_) < 1e-3:
                    return name
            else:
                if torch.dist(param.flatten(), param_.flatten()) < 1e-3:
                    return name
    return None


def build_rename_dict(source_state_dict, target_state_dict, split_qkv=False):
    matched_keys = set()
    with torch.no_grad():
        for name in source_state_dict:
            rename = search_parameter(source_state_dict[name], target_state_dict)
            if rename is not None:
                print(f'"{name}": "{rename}",')
                matched_keys.add(rename)
            elif split_qkv and len(source_state_dict[name].shape)>=1 and source_state_dict[name].shape[0]%3==0:
                length = source_state_dict[name].shape[0] // 3
                rename = []
                for i in range(3):
                    rename.append(search_parameter(source_state_dict[name][i*length: i*length+length], target_state_dict))
                if None not in rename:
                    print(f'"{name}": {rename},')
                    for rename_ in rename:
                        matched_keys.add(rename_)
    for name in target_state_dict:
        if name not in matched_keys:
            print("Cannot find", name, target_state_dict[name].shape)


def search_for_files(folder, extensions):
    files = []
    if os.path.isdir(folder):
        for file in sorted(os.listdir(folder)):
            files += search_for_files(os.path.join(folder, file), extensions)
    elif os.path.isfile(folder):
        for extension in extensions:
            if folder.endswith(extension):
                files.append(folder)
                break
    return files


def convert_state_dict_keys_to_single_str(state_dict, with_shape=True):
    keys = []
    for key, value in state_dict.items():
        if isinstance(key, str):
            if isinstance(value, torch.Tensor):
                if with_shape:
                    shape = "_".join(map(str, list(value.shape)))
                    keys.append(key + ":" + shape)
                keys.append(key)
            elif isinstance(value, dict):
                keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape))
    keys.sort()
    keys_str = ",".join(keys)
    return keys_str


def split_state_dict_with_prefix(state_dict):
    keys = sorted([key for key in state_dict if isinstance(key, str)])
    prefix_dict = {}
    for key in  keys:
        prefix = key if "." not in key else key.split(".")[0]
        if prefix not in prefix_dict:
            prefix_dict[prefix] = []
        prefix_dict[prefix].append(key)
    state_dicts = []
    for prefix, keys in prefix_dict.items():
        sub_state_dict = {key: state_dict[key] for key in keys}
        state_dicts.append(sub_state_dict)
    return state_dicts


def hash_state_dict_keys(state_dict, with_shape=True):
    keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape)
    keys_str = keys_str.encode(encoding="UTF-8")
    return hashlib.md5(keys_str).hexdigest()


import imageio, os, torch, warnings, torchvision, argparse, json
from utils import ModelConfig
from models.utils import load_state_dict
from peft import LoraConfig, inject_adapter_in_model
from PIL import Image
import pandas as pd
from tqdm import tqdm
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs
import numpy as np
import cv2
from bitsandbytes.optim import AdamW8bit
from models.unified_dataset import gen_mask, gen_bbox, gen_points
from utils.eval_depth import test as eval_depth
from utils.eval_normal import test as eval_normal
from utils.eval_matting import test as eval_matting
from lora.flux_lora import FluxLoRALoader

# 找到最后的一个检查点
def find_latest_checkpoint(folder):
    if not os.path.exists(folder):
        return None
    checkpoint_files = [f for f in os.listdir(folder) if f.startswith("step-") and f.endswith(".safetensors")]
    if not checkpoint_files:
        return None
    # 提取步数并找到最大的那个
    latest_checkpoint = max(checkpoint_files, key=lambda x: int(x.split('-')[1].split('.')[0]))
    return os.path.join(folder, latest_checkpoint)

class DiffusionTrainingModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        
        
    def to(self, *args, **kwargs):
        for name, model in self.named_children():
            model.to(*args, **kwargs)
        return self
        
        
    def trainable_modules(self):
        trainable_modules = filter(lambda p: p.requires_grad, self.parameters())
        return trainable_modules
    
    
    def trainable_param_names(self):
        trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.named_parameters()))
        trainable_param_names = set([named_param[0] for named_param in trainable_param_names])
        return trainable_param_names
    
    
    def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None, upcast_dtype=None):
        if lora_alpha is None:
            lora_alpha = lora_rank
        lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules)
        model = inject_adapter_in_model(lora_config, model)
        if upcast_dtype is not None:
            for param in model.parameters():
                if param.requires_grad:
                    param.data = param.to(upcast_dtype)
        return model


    def mapping_lora_state_dict(self, state_dict):
        new_state_dict = {}
        for key, value in state_dict.items():
            if "lora_A.weight" in key or "lora_B.weight" in key:
                new_key = key.replace("lora_A.weight", "lora_A.default.weight").replace("lora_B.weight", "lora_B.default.weight")
                new_state_dict[new_key] = value
            elif "lora_A.default.weight" in key or "lora_B.default.weight" in key:
                new_state_dict[key] = value
        return new_state_dict


    def export_trainable_state_dict(self, state_dict, remove_prefix=None):
        trainable_param_names = self.trainable_param_names()
        state_dict = {name: param for name, param in state_dict.items() if name in trainable_param_names}
        if remove_prefix is not None:
            state_dict_ = {}
            for name, param in state_dict.items():
                if name.startswith(remove_prefix):
                    name = name[len(remove_prefix):]
                state_dict_[name] = param
            state_dict = state_dict_
        return state_dict
    
    
    def transfer_data_to_device(self, data, device):
        for key in data:
            if isinstance(data[key], torch.Tensor):
                data[key] = data[key].to(device)
        return data
    
    
    def parse_model_configs(self, model_paths, model_id_with_origin_paths, enable_fp8_training=False):
        offload_dtype = torch.float8_e4m3fn if enable_fp8_training else None
        model_configs = []
        if model_paths is not None:
            # model_paths = json.loads(model_paths)
            model_paths = model_paths.split(",")
            model_configs += [ModelConfig(path=path, offload_dtype=offload_dtype) for path in model_paths]
        if model_id_with_origin_paths is not None:
            model_id_with_origin_paths = model_id_with_origin_paths.split(",")
            model_configs += [ModelConfig(model_id=i.split(":")[0], origin_file_pattern=i.split(":")[1], offload_dtype=offload_dtype) for i in model_id_with_origin_paths]
        return model_configs


    def switch_pipe_to_training_mode(
        self,
        pipe,
        trainable_models,
        lora_base_model=None, lora_target_modules=None, lora_rank=None, lora_checkpoint=None,
        enable_fp8_training=False,
    ):
        # Scheduler
        pipe.scheduler.set_timesteps(1000, training=True)
        
        # Freeze untrainable models
        pipe.freeze_except([] if trainable_models is None else trainable_models.split(","))
        
        # Enable FP8 if pipeline supports
        if enable_fp8_training and hasattr(pipe, "_enable_fp8_lora_training"):
            pipe._enable_fp8_lora_training(torch.float8_e4m3fn)
        
        # Add LoRA to the base models
        if lora_base_model is not None:
            model = self.add_lora_to_model(
                getattr(pipe, lora_base_model),
                target_modules=lora_target_modules.split(","),
                lora_rank=lora_rank,
                upcast_dtype=pipe.torch_dtype,
            )
            if lora_checkpoint is not None:
                state_dict = load_state_dict(lora_checkpoint)
                loader = FluxLoRALoader(torch_dtype=torch.bfloat16, device="cuda")
                state_dict = loader.convert_state_dict(state_dict)
                state_dict = self.mapping_lora_state_dict(state_dict)
                load_result = model.load_state_dict(state_dict, strict=False)
                print(f"LoRA checkpoint loaded: {lora_checkpoint}, total {len(state_dict)} keys")
                if len(load_result[1]) > 0:
                    print(f"Warning, LoRA key mismatch! Unexpected keys in LoRA checkpoint: {load_result[1]}")
            setattr(pipe, lora_base_model, model)


class ModelLogger:
    def __init__(self, output_path, remove_prefix_in_ckpt=None, state_dict_converter=lambda x:x,
                 args=None):
        self.output_path = output_path
        self.remove_prefix_in_ckpt = remove_prefix_in_ckpt
        self.state_dict_converter = state_dict_converter
        self.num_steps = args.resume_steps
        # Evaluation related
        self.eval_steps = args.eval_steps
        self.eval_file_list = args.eval_file_list
        self.eval_output_dir = os.path.join(self.output_path, "eval")
        self.eval_prompt = args.default_caption
        self.eval_num_inference_steps = args.eval_num_inference_steps
        self.eval_embedded_guidance = args.eval_embedded_guidance
        self.eval_height = args.height
        self.eval_width = args.width
        self.trainable_models = args.trainable_models
        self.using_log = args.using_log
        self.using_disp = args.using_disp
        self.using_sqrt = args.using_sqrt
        self.using_sqrt_disp = args.using_sqrt_disp
        self.task = args.task
        self.deterministic_flow = args.deterministic_flow
        self.eval_args = argparse.Namespace(
            pred_path="result/nyu_depth_v2",  # 对应 --pred_path
            gt_path="../dataset/Eval/depth/nyuv2",  # 对应 --gt_path
            dataset="nyu",  # 目标脚本默认值（命令行没传，用默认）
            eigen_crop=False,  
            garg_crop=False,  
            min_depth_eval=1e-3, 
            max_depth_eval=10.0, 
            do_kb_crop=False,
            no_verbose=False,
            using_log=args.using_log,
            using_disp=args.using_disp,
            using_sqrt=args.using_sqrt,
            using_sqrt_disp=args.using_sqrt_disp,
            using_pdf=args.using_pdf,
        )
        self.matting_prompt = args.matting_prompt
        # lora related
        self.lora_base_model = args.lora_base_model
        self.lora_target_modules = args.lora_target_modules
        self.lora_rank = args.lora_rank
        self.lora_checkpoint = args.lora_checkpoint

    def on_step_end(self, accelerator, model, save_steps=None):
        self.num_steps += 1
        if save_steps is not None and self.num_steps % save_steps == 0:
            self.save_model(accelerator, model, f"step-{self.num_steps}.safetensors")


    def run_evaluation(self, model, full_size=False):
        if self.eval_steps is None:
            return
        if not full_size:
            files = self.eval_file_list[:10]
            output_dir = self.eval_output_dir
        else:
            files = self.eval_file_list[:654]
            output_dir = self.eval_output_dir + "_full"
        if len(files) == 0:
            return

        try:
            from models.unified_dataset import UnifiedDataset
            transform = UnifiedDataset.default_image_operator()
        except Exception as e:
            print(f"[ModelLogger][Eval] Failed to build transform: {e}")
            return
        pipe = getattr(model, 'pipe', None)
        if pipe is None:
            print("[ModelLogger][Eval] Model has no 'pipe' attribute; skipping evaluation.")
            return
        # Maintain relative folder structure
        print(f"[ModelLogger][Eval] Running evaluation on {len(files)} files (steps={self.num_steps}). Saving to {output_dir}")
        if self.task == "depth" or self.task == "normal":
            base_dir = f"/mnt/nfs/workspace/syq/dataset/Eval/{self.task}/nyuv2"
        elif self.task == "matting":
            base_dir = "/mnt/nfs/workspace/syq/dataset/matting/P3M-10k"
        else:
            raise ValueError(f"Unknown task {self.task}")
        os.makedirs(output_dir, exist_ok=True)
        for file in files:
            try:
                save_to = file.replace(base_dir, output_dir).replace(".png", ".npy").replace(".jpg", ".npy")
                os.makedirs(os.path.dirname(save_to), exist_ok=True)
                with torch.no_grad():
                    if "kontext" in output_dir:
                        if self.task == "depth" or self.task == "normal":
                            output = pipe(
                                prompt=self.eval_prompt,
                                kontext_images=transform(file),
                                height=self.eval_height,
                                width=self.eval_width,
                                embedded_guidance=self.eval_embedded_guidance,
                                num_inference_steps=self.eval_num_inference_steps,
                                seed=42,
                                output_type="np",
                                rand_device=pipe.device,  # ensure format
                                deterministic_flow=self.deterministic_flow,
                                task=self.task,
                            )
                        elif self.task == "matting":
                            alpha = np.array(Image.open(file.replace("blurred_image","mask").replace("original_image","mask").replace(".jpg",".png")).convert("L")) / 255.0
                            if self.matting_prompt is not None:
                                if self.matting_prompt == "trimap":
                                    visual_prompt = transform(file.replace("blurred_image","trimap").replace("original_image","trimap").replace(".jpg",".png"))
                                elif self.matting_prompt == "mask":
                                    visual_prompt,visual_prompt_coords = gen_mask(alpha)
                                elif self.matting_prompt == "bbox":
                                    visual_prompt,visual_prompt_coords = gen_bbox(alpha,0)
                                elif self.matting_prompt == "points":
                                    visual_prompt,visual_prompt_coords = gen_points(alpha,radius=30)
                                
                                if isinstance(visual_prompt, np.ndarray):
                                    # resize to 1/8 size
                                    # visual_prompt = cv2.resize(visual_prompt, (self.eval_width // 8, self.eval_height // 8), interpolation=cv2.INTER_NEAREST)
                                    # visual_prompt = torch.from_numpy(visual_prompt).unsqueeze(0).to(torch.bfloat16).to("cuda")
                                    # visual_prompt_coords = torch.from_numpy(visual_prompt_coords).to(torch.bfloat16).to("cuda")
                                    # del alpha
                                    visual_prompt = cv2.resize(visual_prompt,(self.eval_width,self.eval_height),interpolation=cv2.INTER_LINEAR)
                                    visual_prompt = torch.from_numpy(visual_prompt*2 - 1).repeat(3,1,1).to(pipe.device).to(torch.bfloat16)
                                kontext_images = [transform(file),visual_prompt]
                            output = pipe(
                                prompt=self.eval_prompt,
                                kontext_images=kontext_images,
                                height=self.eval_height,
                                width=self.eval_width,
                                cfg_scale=self.eval_embedded_guidance,
                                num_inference_steps=self.eval_num_inference_steps,
                                seed=42,
                                output_type="np",
                                rand_device=pipe.device,  # ensure format
                                deterministic_flow=self.deterministic_flow,
                                task=self.task,
                                # kontext_masks=visual_prompt if self.matting_prompt is not None else None,
                            )
                    elif "qwen" in output_dir:
                        output = pipe(
                            prompt=self.eval_prompt,
                            edit_image=transform(file),
                            height=self.eval_height,
                            width=self.eval_width,
                            cfg_scale=self.eval_embedded_guidance,
                            num_inference_steps=self.eval_num_inference_steps,
                            seed=42,
                            output_type="np",
                            rand_device=pipe.device,  # ensure format
                            # deterministic_flow=self.deterministic_flow,
                            task=self.task,
                        )
                np.save(save_to, output)
            except Exception as e:
                print(f"[ModelLogger][Eval] Failed on file {file}: {e}")
        print("[ModelLogger][Eval] Inference finished.")
        self.eval_args.pred_path = output_dir
        if self.task == "depth":
            return eval_depth(self.eval_args)
        elif self.task == "normal":
            self.eval_args.gt_path = "/mnt/nfs/workspace/syq/dataset/Eval/normal/nyuv2"
            return eval_normal(self.eval_args)
        elif self.task == "matting":
            self.eval_args.gt_path = "/mnt/nfs/workspace/syq/dataset/matting/P3M-10k"
            self.eval_args.dataset = "p3m-np"
            return eval_matting(self.eval_args)
        else:
            print(f"[ModelLogger][Eval] Unknown task {self.task}; skipping evaluation.")
            return

    def on_eval_step(self, accelerator, model: DiffusionTrainingModule):
        if self.eval_steps is None or self.num_steps % self.eval_steps != 0:
            return
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            results = self.run_evaluation(model)
            with open(os.path.join(self.eval_output_dir, "log.txt"), "a") as f:
                f.write(f"Step {self.num_steps} evaluation results:\n")
                f.write(results+"\n")
        accelerator.wait_for_everyone()
        
        model.switch_pipe_to_training_mode(model.pipe,self.trainable_models,None,None,None,None)


    def on_epoch_end(self, accelerator, model, epoch_id):
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            state_dict = accelerator.get_state_dict(model)
            state_dict = accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix=self.remove_prefix_in_ckpt)
            state_dict = self.state_dict_converter(state_dict)
            os.makedirs(self.output_path, exist_ok=True)
            path = os.path.join(self.output_path, f"epoch-{epoch_id}.safetensors")
            accelerator.save(state_dict, path, safe_serialization=True)


    def on_training_end(self, accelerator, model, save_steps=None):
        if save_steps is not None and self.num_steps % save_steps != 0:
            self.save_model(accelerator, model, f"step-{self.num_steps}.safetensors")


    def save_model(self, accelerator, model, file_name):
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            state_dict = accelerator.get_state_dict(model)
            state_dict = accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix=self.remove_prefix_in_ckpt)
            state_dict = self.state_dict_converter(state_dict)
            os.makedirs(self.output_path, exist_ok=True)
            path = os.path.join(self.output_path, file_name)
            accelerator.save(state_dict, path, safe_serialization=True)
            results = self.run_evaluation(model,full_size=True)
            with open(os.path.join(self.eval_output_dir+"_full", "log_full.txt"), "a") as f:
                f.write(f"Step {self.num_steps} evaluation results:\n")
                f.write(results+"\n")
        accelerator.wait_for_everyone()

        model.switch_pipe_to_training_mode(model.pipe,self.trainable_models,None,None,None,None)
def launch_training_task(
    dataset: torch.utils.data.Dataset,
    model: DiffusionTrainingModule,
    model_logger: ModelLogger,
    dataset_sampler: torch.utils.data.Sampler = None,
    learning_rate: float = 1e-5,
    weight_decay: float = 1e-2,
    num_workers: int = 8,
    save_steps: int = None,
    num_epochs: int = 1,
    gradient_accumulation_steps: int = 1,
    find_unused_parameters: bool = False,
    args = None,
):
    if args is not None:
        learning_rate = args.learning_rate
        weight_decay = args.weight_decay
        num_workers = args.dataset_num_workers
        save_steps = args.save_steps
        num_epochs = args.num_epochs
        gradient_accumulation_steps = args.gradient_accumulation_steps
        find_unused_parameters = args.find_unused_parameters
    
    if args.adamw8bit:
        print("Using 8-bit AdamW optimizer.")
        optimizer = AdamW8bit(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay)
    else:
        print("Using regular AdamW optimizer.") 
        optimizer = torch.optim.AdamW(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay)

    scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
    # dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=num_workers)
    if dataset_sampler is None:
        # random batch sampler with batch_size = args.batch_size
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=num_workers)
    else:
        dataloader = torch.utils.data.DataLoader(dataset, batch_sampler=dataset_sampler, num_workers=num_workers)
    

    accelerator = Accelerator(
        gradient_accumulation_steps=gradient_accumulation_steps,
        kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=find_unused_parameters)],
    )
    if getattr(accelerator.state, "deepspeed_plugin", None) is not None:
        print("Using DeepSpeed with batch size on per GPU", args.batch_size)
        accelerator.state.deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = int(args.batch_size)
    model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
    
    for epoch_id in range(num_epochs):
        loss = None
        pbar = tqdm(dataloader, desc=f"Epoch {epoch_id}")
        for i,data in enumerate(pbar):
            with accelerator.accumulate(model):
                optimizer.zero_grad()
                # if dataset.load_from_cache:
                #     loss = model({}, inputs=data)
                # else:
                    # loss = model(data)
                loss = model(data)
                # if i%50==0 and accelerator.is_main_process:
                #     with open("debug_loss_uni_flux.txt","a") as f:
                #         f.write(f"step {i}: {loss.item()}\n")
                if isinstance(loss, tuple) or isinstance(loss, list):
                    if epoch_id >= args.extra_loss_start_epoch:
                        loss, extra_loss = loss[0], loss[1]
                        coeff_scale = loss.detach().cpu().item() / (extra_loss.detach().cpu().item()+1e-3)
                        coeff_step = max(0,epoch_id-args.extra_loss_start_epoch+i/len(dataloader))
                        loss = loss + extra_loss * coeff_scale * coeff_step
                    else:
                        loss = loss[0]
                pbar.set_description(f"Loss {loss.item():.4f}")
                accelerator.backward(loss)
                optimizer.step()
                torch.cuda.empty_cache()
                model_logger.on_step_end(accelerator, model, save_steps)
                # Evaluation hook
                model_logger.on_eval_step(accelerator, model)
                scheduler.step()
        if save_steps is None:
            model_logger.on_epoch_end(accelerator, model, epoch_id)
    model_logger.on_training_end(accelerator, model, save_steps)

def parse_flux_model_configs(root_path):
    # given the root path, and then load the following: 
    # text_encoder, text_encoder_2, tokenizer, tokenizer_2, transformer(a folder of 3 parts) / or flux1-kontext-dev.safetensors, vae
    model_configs = []
    _targets = ["flux1-kontext-dev.safetensors", "text_encoder/model.safetensors", "text_encoder_2","ae.safetensors"]
    if "kontext" not in root_path.lower():
        _targets = ["flux1-dev.safetensors", "text_encoder/model.safetensors", "text_encoder_2","ae.safetensors"]
    for model_name in _targets:
        model_path = os.path.join(root_path, model_name)
        model_configs.append(ModelConfig(path=model_path))
    return model_configs


def flux_parser():
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument("--dataset_base_path", type=str, default="", required=True, help="Base path of the dataset.")
    parser.add_argument("--dataset_metadata_path", type=str, default=None, help="Path to the metadata file of the dataset.")
    parser.add_argument("--max_pixels", type=int, default=1024*1024, help="Maximum number of pixels per frame, used for dynamic resolution..")
    parser.add_argument("--height", type=int, default=768, help="Height of images. Leave `height` and `width` empty to enable dynamic resolution.")
    parser.add_argument("--width", type=int, default=768, help="Width of images. Leave `height` and `width` empty to enable dynamic resolution.")
    parser.add_argument("--data_file_keys", type=str, default="image", help="Data file keys in the metadata. Comma-separated.")
    parser.add_argument("--dataset_repeat", type=int, default=1, help="Number of times to repeat the dataset per epoch.")
    parser.add_argument("--model_paths", type=str, default=None, help="Paths to load models. In JSON format.")
    parser.add_argument("--model_id_with_origin_paths", type=str, default=None, help="Model ID with origin paths, e.g., Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors. Comma-separated.")
    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate.")
    parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs.")
    parser.add_argument("--output_path", type=str, default="./models", help="Output save path.")
    parser.add_argument("--remove_prefix_in_ckpt", type=str, default="pipe.dit.", help="Remove prefix in ckpt.")
    parser.add_argument("--trainable_models", type=str, default=None, help="Models to train, e.g., dit, vae, text_encoder.")
    parser.add_argument("--lora_base_model", type=str, default=None, help="Which model LoRA is added to.")
    parser.add_argument("--lora_target_modules", type=str, default="q,k,v,o,ffn.0,ffn.2", help="Which layers LoRA is added to.")
    parser.add_argument("--lora_rank", type=int, default=32, help="Rank of LoRA.")
    parser.add_argument("--lora_checkpoint", type=str, default=None, help="Path to the LoRA checkpoint. If provided, LoRA will be loaded from this checkpoint.")
    parser.add_argument("--extra_inputs", default=None, help="Additional model inputs, comma-separated.")
    parser.add_argument("--align_to_opensource_format", default=False, action="store_true", help="Whether to align the lora format to opensource format. Only for DiT's LoRA.")
    parser.add_argument("--use_gradient_checkpointing", default=False, action="store_true", help="Whether to use gradient checkpointing.")
    parser.add_argument("--use_gradient_checkpointing_offload", default=False, action="store_true", help="Whether to offload gradient checkpointing to CPU memory.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps.")
    parser.add_argument("--find_unused_parameters", default=False, action="store_true", help="Whether to find unused parameters in DDP.")
    parser.add_argument("--save_steps", type=int, default=None, help="Number of checkpoint saving invervals. If None, checkpoints will be saved every epoch.")
    parser.add_argument("--dataset_num_workers", type=int, default=16, help="Number of workers for data loading.")
    parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay.")
    parser.add_argument("--mixed_sampler", default=False, action="store_true", help="Whether to use mixed sampler (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
    # normalization
    parser.add_argument("--depth_normalization",type=str, default=None, help="Normalization method for depth map")
    parser.add_argument("--using_log", default=False, action="store_true", help="Whether to use log for depth preprocessing.")
    parser.add_argument("--using_disp", default=False, action="store_true", help="Whether to use disp for depth preprocessing.")
    parser.add_argument("--using_sqrt",default=False, action="store_true", help="Whether to use sqrt for depth preprocessing.")
    parser.add_argument("--using_sqrt_disp",default=False, action="store_true", help="Whether to use sqrt for depth preprocessing.")
    parser.add_argument("--using_pdf", default=False, action="store_true", help="Whether to use pdf for depth preprocessing.")
    # text and visual prompt
    parser.add_argument("--with_mask", default=False, action="store_true", help="Whether to use mask for loss calculation, espcially for normal estimation.")
    parser.add_argument("--default_caption", type=str, default=None, help="Default caption for all training samples.")
    parser.add_argument("--matting_prompt", type=str, default=None,choices=["trimap", "mask", "bbox", "points"], help="Prompt for image matting.")
    parser.add_argument("--use_coor_input", default=False, action="store_true", help="Whether to use coordinate as model input.")
    parser.add_argument("--use_camera_intrinsics", default=False, action="store_true", help="Whether to use camera intrinsics as model input.")
    # Evaluation related arguments
    parser.add_argument("--eval_steps", type=int, default=5, help="Run evaluation every N training steps.")
    parser.add_argument("--eval_num_inference_steps", type=int, default=1, help="Number of inference steps for evaluation.")
    parser.add_argument("--eval_embedded_guidance", type=float, default=1, help="Embedded guidance for evaluation generation.")
    parser.add_argument("--eval_file_list", type=str, default="", help="A text file containing a list of file paths to be used for evaluation. If empty, no evaluation is performed.")
    # deterministic training
    parser.add_argument("--deterministic_flow", default=False, action="store_true", help="Whether to use deterministic flow for training.")
    # training
    parser.add_argument("--batch_size", type=int, default=1, help="Batch size for training.")
    parser.add_argument("--adamw8bit", default=False, action="store_true", help="Whether to use 8-bit Adam optimizer.")
    parser.add_argument("--multi_res_noise", default=False, action="store_true", help="Whether to use multi-resolution noise for training.")
    parser.add_argument("--resume", default=False, action="store_true", help="Whether to resume training from a checkpoint.")
    parser.add_argument("--task", type=str, default="depth", required=False, help="Task type, e.g., depth, normal.")
    # loss
    parser.add_argument("--extra_loss", type=str, default=None, help="Loss type for depth estimation, e.g., l1, l2, berhu.")
    parser.add_argument("--extra_loss_start_epoch", type=int, default=0, help="The start epoch for applying extra loss. before this epoch, only the main loss is used.")
    return parser