combined_script.py

import os
import toml

def process_flux_checkpoint(project_root):
    print("[INFO] Starting FluxCheckpoint configuration...")
    
    # ------------------------------
    # Initial Setup for FluxCheckpoint
    # ------------------------------
    output_dir = os.path.join(project_root, "output")
    batch_config_dir = os.path.join(project_root, "BatchConfig", "Flux")
    os.makedirs(batch_config_dir, exist_ok=True)
    
    # Base Configuration for FluxCheckpoint
    base_config_flux = """
adaptive_noise_scale = 0
ae = "./models/trainX/ae.safetensors"
blocks_to_swap = 9
bucket_no_upscale = true
bucket_reso_steps = 64
cache_latents = true
cache_latents_to_disk = true
cache_text_encoder_outputs = true
cache_text_encoder_outputs_to_disk = true
caption_dropout_every_n_epochs = 0
caption_dropout_rate = 0
caption_extension = ".txt"
clip_l = "./models/trainX/clip_l.safetensors"
discrete_flow_shift = 3.1582
double_blocks_to_swap = 0
dynamo_backend = "no"
max_train_epochs = 1
full_bf16 = true
fused_backward_pass = true
gradient_accumulation_steps = 1
gradient_checkpointing = true
guidance_scale = 1
huber_c = 0.1
huber_schedule = "snr"
keep_tokens = 0
learning_rate = 4e-6
learning_rate_te = 0
logging_dir = "Z"
loss_type = "l2"
lr_scheduler = "constant"
lr_scheduler_args = []
lr_scheduler_num_cycles = 1
lr_scheduler_power = 1
lr_warmup_steps = 0
max_bucket_reso = 2048
max_data_loader_n_workers = 0
max_timestep = 1000
max_token_length = 75
mem_eff_save = true
min_bucket_reso = 256
mixed_precision = "bf16"
model_prediction_type = "raw"
multires_noise_discount = 0.3
multires_noise_iterations = 0
noise_offset = 0
noise_offset_type = "Original"
optimizer_args = [ "scale_parameter=False", "relative_step=False", "warmup_init=False", "weight_decay=0.01",]
optimizer_type = "Adafactor"
output_dir = "Z"
output_name = "m1u"
persistent_data_loader_workers = 0
pretrained_model_name_or_path = "./models/trainX/flux1-dev-fp8.safetensors"
prior_loss_weight = 1
resolution = "1024,1024"
sample_every_n_steps = 250
sample_prompts = "./sample_prompts.txt"
sample_sampler = "euler_a"
save_model_as = "safetensors"
save_precision = "fp16"
seed = 1
single_blocks_to_swap = 0
t5xxl = "./models/trainX/t5xxl_fp8_e4m3fn.safetensors"
t5xxl_max_token_length = 512
timestep_sampling = "sigmoid"
train_batch_size = 1
train_blocks = "all"
train_data_dir = "Z"
vae_batch_size = 4
wandb_run_name = "m1u"
xformers = true
"""
    try:
        config_dict_base_flux = toml.loads(base_config_flux)
    except toml.TomlDecodeError as e:
        print(f"[ERROR] Failed to parse FluxCheckpoint base configuration: {e}")
        return False

    # Get list of folders in output directory
    try:
        folders = [f for f in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, f))]
    except FileNotFoundError:
        print(f"[WARNING] The directory '{output_dir}' does not exist.")
        folders = []

    if not folders:
        print(f"[INFO] No folders found in '{output_dir}'.")
    else:
        for folder in folders:
            # Create necessary subdirectories
            subdirs = ["model", "log", "img"]
            for subdir in subdirs:
                subdir_path = os.path.join(output_dir, folder, subdir)
                os.makedirs(subdir_path, exist_ok=True)

            # Create a copy of the base configuration
            config_dict = config_dict_base_flux.copy()

            # Define keys that represent paths
            path_keys = [
                "ae",
                "clip_l",
                "pretrained_model_name_or_path",
                "t5xxl",
                "sample_prompts"
            ]

            # Convert relative paths to absolute paths with forward slashes
            for key in path_keys:
                if key in config_dict:
                    original_path = config_dict[key]
                    absolute_path = os.path.abspath(os.path.join(project_root, original_path))
                    absolute_path = absolute_path.replace('\\', '/')
                    config_dict[key] = absolute_path

            # Update specific paths based on folder name
            config_dict['output_dir'] = os.path.abspath(os.path.join(output_dir, folder, "model")).replace('\\', '/')
            config_dict['output_name'] = folder
            config_dict['logging_dir'] = os.path.abspath(os.path.join(output_dir, folder, "log")).replace('\\', '/')
            config_dict['train_data_dir'] = os.path.abspath(os.path.join(output_dir, folder, "img")).replace('\\', '/')

            # Verify the existence of paths
            required_paths = [
                config_dict.get("ae"),
                config_dict.get("clip_l"),
                config_dict.get("pretrained_model_name_or_path"),
                config_dict.get("t5xxl"),
                config_dict.get("sample_prompts"),
                config_dict.get("output_dir"),
                config_dict.get("logging_dir"),
                config_dict.get("train_data_dir")
            ]

            missing_paths = [path for path in required_paths if path and not os.path.exists(path)]
            if missing_paths:
                print(f"[WARNING] Some paths do not exist for folder '{folder}':")
                for path in missing_paths:
                    print(f"  - {path}")
                print("Please ensure that all necessary files and directories are present.\n")

            # Create the TOML file name and path
            toml_filename = f"{folder}.toml"
            toml_path = os.path.join(batch_config_dir, toml_filename)

            # Write the configuration in TOML format
            try:
                with open(toml_path, 'w') as toml_file:
                    toml.dump(config_dict, toml_file)
                print(f"[SUCCESS] FluxCheckpoint configuration file created: {toml_path}")
            except Exception as e:
                print(f"[ERROR] Failed to write FluxCheckpoint TOML file '{toml_path}': {e}")

    print("[INFO] FluxCheckpoint configuration completed.\n")
    return True

def process_flux_lora(project_root):
    print("[INFO] Starting FluxLora configuration...")
    
    # ------------------------------
    # Initial Setup for FluxLora
    # ------------------------------
    output_dir = os.path.join(project_root, "output")
    batch_config_dir = os.path.join(project_root, "BatchConfig", "FluxLORA")
    os.makedirs(batch_config_dir, exist_ok=True)
    
    # Base Configuration for FluxLora
    base_config_lora = """
ae = "./models/trainX/ae.safetensors"
bucket_no_upscale = true
bucket_reso_steps = 64
cache_latents = true
cache_latents_to_disk = true
cache_text_encoder_outputs = true
cache_text_encoder_outputs_to_disk = true
caption_extension = ".txt"
clip_l = "./models/trainX/clip_l.safetensors"
discrete_flow_shift = 3.1582
dynamo_backend = "no"
max_train_epochs = 1
fp8_base = true
full_bf16 = true
gradient_accumulation_steps = 1
gradient_checkpointing = true
guidance_scale = 1.0
highvram = true
huber_c = 0.1
huber_schedule = "snr"
logging_dir = "X"
loss_type = "l2"
lr_scheduler = "constant"
lr_scheduler_args = []
lr_scheduler_num_cycles = 1
lr_scheduler_power = 1
max_bucket_reso = 2048
max_data_loader_n_workers = 0
max_timestep = 1000
mem_eff_save = true
min_bucket_reso = 256
mixed_precision = "bf16"
model_prediction_type = "raw"
network_alpha = 128
network_args = ["train_double_block_indices=all", "train_single_block_indices=all"]
network_dim = 128
network_module = "networks.lora_flux"
noise_offset_type = "Original"
optimizer_args = ["scale_parameter=False", "relative_step=False", "warmup_init=False", "weight_decay=0.01"]
optimizer_type = "Adafactor"
output_dir = "X"
output_name = "X"
pretrained_model_name_or_path = "./models/trainX/flux1-dev-fp8.safetensors"
prior_loss_weight = 1
resolution = "1024,1280"
sample_every_n_steps = 250
sample_prompts = "./sample_prompts.txt"
sample_sampler = "euler_a"
save_model_as = "safetensors"
save_precision = "float"
seed = 1
t5xxl = "./models/trainX/t5xxl_fp8_e4m3fn.safetensors"
t5xxl_max_token_length = 512
text_encoder_lr = [5e-5, 5e-5]
timestep_sampling = "sigmoid"
train_batch_size = 1
train_data_dir = "X"
unet_lr = 5e-5
vae_batch_size = 4
wandb_run_name = "X"
xformers = true
"""
    try:
        config_dict_base_lora = toml.loads(base_config_lora)
    except toml.TomlDecodeError as e:
        print(f"[ERROR] Failed to parse FluxLora base configuration: {e}")
        return False

    # Get list of folders in output directory
    try:
        folders = [f for f in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, f))]
    except FileNotFoundError:
        print(f"[WARNING] The directory '{output_dir}' does not exist.")
        folders = []

    if not folders:
        print(f"[INFO] No folders found in '{output_dir}'.")
    else:
        for folder in folders:
            # Create necessary subdirectories
            subdirs = ["model", "log", "img"]
            for subdir in subdirs:
                subdir_path = os.path.join(output_dir, folder, subdir)
                os.makedirs(subdir_path, exist_ok=True)

            # Create a copy of the base configuration
            config_dict = config_dict_base_lora.copy()

            # Define keys that represent paths
            path_keys = [
                "ae",
                "clip_l",
                "pretrained_model_name_or_path",
                "t5xxl",
                "sample_prompts",
                "output_dir",
                "logging_dir",
                "train_data_dir"
            ]

            # Convert relative paths to absolute paths with forward slashes
            for key in path_keys:
                if key in config_dict and config_dict[key] != "X":
                    original_path = config_dict[key]
                    absolute_path = os.path.abspath(os.path.join(project_root, original_path))
                    absolute_path = absolute_path.replace('\\', '/')
                    config_dict[key] = absolute_path
                elif key in config_dict and config_dict[key] == "X":
                    # Update paths based on the current folder
                    if key == "output_dir":
                        config_dict[key] = os.path.abspath(os.path.join(output_dir, folder, "model")).replace('\\', '/')
                    elif key == "logging_dir":
                        config_dict[key] = os.path.abspath(os.path.join(output_dir, folder, "log")).replace('\\', '/')
                    elif key == "train_data_dir":
                        config_dict[key] = os.path.abspath(os.path.join(output_dir, folder, "img")).replace('\\', '/')
                    else:
                        # For keys that are set to "X" and not explicitly handled
                        config_dict[key] = "X"

            # Update 'output_name' based on folder name
            config_dict['output_name'] = folder

            # Verify the existence of paths
            required_paths = [
                config_dict.get("ae"),
                config_dict.get("clip_l"),
                config_dict.get("pretrained_model_name_or_path"),
                config_dict.get("t5xxl"),
                config_dict.get("sample_prompts"),
                config_dict.get("output_dir"),
                config_dict.get("logging_dir"),
                config_dict.get("train_data_dir")
            ]

            missing_paths = [path for path in required_paths if path and path != "X" and not os.path.exists(path)]
            if missing_paths:
                print(f"[WARNING] Some paths do not exist for folder '{folder}':")
                for path in missing_paths:
                    print(f"  - {path}")
                print("Please ensure that all necessary files and directories are present.\n")
            else:
                print(f"[INFO] All required paths exist for folder '{folder}'.")

            # Create the TOML file name and path
            toml_filename = f"{folder}.toml"
            toml_path = os.path.join(batch_config_dir, toml_filename)

            # Write the configuration in TOML format
            try:
                with open(toml_path, 'w') as toml_file:
                    toml.dump(config_dict, toml_file)
                print(f"[SUCCESS] FluxLora configuration file created: {toml_path}")
            except Exception as e:
                print(f"[ERROR] Failed to write FluxLora TOML file '{toml_path}': {e}")

    print("[INFO] FluxLora configuration completed.\n")
    return True

def process_sdxl_nude(project_root):
    print("[INFO] Starting SDXLNude configuration...")
    
    # ------------------------------
    # Initial Setup for SDXLNude
    # ------------------------------
    output_dir = os.path.join(project_root, "output")
    batch_config_dir = os.path.join(project_root, "BatchConfig", "Nude")
    os.makedirs(batch_config_dir, exist_ok=True)
    
    sample_prompts_path = os.path.join(project_root, "sample_prompts.txt")
    
    # Base Configuration for SDXLNude
    base_config_nude = """
adaptive_noise_scale = 0
bucket_no_upscale = true
bucket_reso_steps = 64
cache_latents = true
cache_latents_to_disk = true
caption_dropout_every_n_epochs = 0
caption_dropout_rate = 0
clip_skip = 1
dynamo_backend = "no"
max_train_epochs = 10
full_bf16 = true
gradient_accumulation_steps = 1
gradient_checkpointing = true
huber_c = 0.1
huber_schedule = "snr"
ip_noise_gamma = 0.1
keep_tokens = 0
learning_rate = 8e-6
learning_rate_te1 = 3e-6
learning_rate_te2 = 0
logging_dir = "C:/Users/newec/Documents/sepuedeborrar"
loss_type = "l2"
lr_scheduler = "constant"
lr_scheduler_args = []
lr_scheduler_num_cycles = 1
lr_scheduler_power = 1
lr_warmup_steps = 0
max_bucket_reso = 2048
max_data_loader_n_workers = 0
max_timestep = 1000
max_token_length = 75
min_bucket_reso = 256
mixed_precision = "bf16"
multires_noise_discount = 0
multires_noise_iterations = 0
noise_offset = 0.0357
noise_offset_type = "Original"
optimizer_args = [ "scale_parameter=False", "relative_step=False", "warmup_init=False",
"weight_decay=0.01",]
optimizer_type = "Adafactor"
output_dir = "C:/Users/newec/Documents/sepuedeborrar"
output_name = "indaX"
persistent_data_loader_workers = 0
pretrained_model_name_or_path = "models/NudeSDXLModel.safetensors"
prior_loss_weight = 1
resolution = "1024,1280"
sample_every_n_steps = 250
sample_prompts = ""
sample_sampler = "euler_a"
save_model_as = "safetensors"
save_precision = "bf16"
train_batch_size = 1
train_data_dir = "C:/Users/newec/Documents/sepuedeborrar"
train_text_encoder = true
vae = "stabilityai/sdxl-vae"
vae_batch_size = 2
wandb_run_name = "X"
"""
    try:
        config_dict_base_nude = toml.loads(base_config_nude)
    except toml.TomlDecodeError as e:
        print(f"[ERROR] Failed to parse SDXLNude base configuration: {e}")
        return False

    # Update the 'sample_prompts' key with the absolute path
    config_dict_base_nude["sample_prompts"] = sample_prompts_path.replace('\\', '/')

    # Get list of folders in output directory
    try:
        folders = [f for f in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, f))]
    except FileNotFoundError:
        print(f"[WARNING] The directory '{output_dir}' does not exist.")
        folders = []

    if not folders:
        print(f"[INFO] No folders found in '{output_dir}'.")
    else:
        for folder in folders:
            # Create necessary subdirectories
            subdirs = ["model", "log", "img"]
            for subdir in subdirs:
                subdir_path = os.path.join(output_dir, folder, subdir)
                os.makedirs(subdir_path, exist_ok=True)
                print(f"[INFO] Subdirectory '{subdir_path}' ensured.")

            # Create a copy of the base configuration
            config_dict = config_dict_base_nude.copy()

            # Define keys that represent paths
            path_keys = [
                "pretrained_model_name_or_path"
            ]

            # Convert relative paths to absolute paths with forward slashes
            for key in path_keys:
                if key in config_dict:
                    original_path = config_dict[key]
                    absolute_path = os.path.abspath(os.path.join(project_root, "models/trainX/NudeSDXLModel.safetensors"))
                    absolute_path = absolute_path.replace('\\', '/')
                    config_dict[key] = absolute_path

            # Update specific paths based on folder name
            config_dict['output_dir'] = os.path.abspath(os.path.join(output_dir, folder, "model")).replace('\\', '/')
            config_dict['output_name'] = folder
            config_dict['logging_dir'] = os.path.abspath(os.path.join(output_dir, folder, "log")).replace('\\', '/')
            config_dict['train_data_dir'] = os.path.abspath(os.path.join(output_dir, folder, "img")).replace('\\', '/')

            # Verify the existence of paths
            required_paths = [
                config_dict.get("pretrained_model_name_or_path"),
                config_dict.get("sample_prompts"),
                config_dict.get("output_dir"),
                config_dict.get("logging_dir"),
                config_dict.get("train_data_dir")
            ]

            missing_paths = [path for path in required_paths if path and not os.path.exists(path)]
            if missing_paths:
                print(f"[WARNING] Some paths do not exist for folder '{folder}':")
                for path in missing_paths:
                    print(f"  - {path}")
                print("Please ensure that all necessary files and directories are present.\n")
            else:
                print(f"[INFO] All required paths exist for folder '{folder}'.")

            # Create the TOML file name and path
            toml_filename = f"{folder}.toml"
            toml_path = os.path.join(batch_config_dir, toml_filename)

            # Write the configuration in TOML format
            try:
                with open(toml_path, 'w') as toml_file:
                    toml.dump(config_dict, toml_file)
                print(f"[SUCCESS] SDXLNude configuration file created: {toml_path}\n")
            except Exception as e:
                print(f"[ERROR] Failed to write SDXLNude TOML file '{toml_path}': {e}\n")

        print("[INFO] All SDXLNude configuration files have been generated successfully.\n")

    # Verify required subdirectories
    if folders:
        all_missing = False
        for folder in folders:
            required_subdirs = [
                os.path.join(output_dir, folder, "model"),
                os.path.join(output_dir, folder, "log"),
                os.path.join(output_dir, folder, "img")
            ]

            for path in required_subdirs:
                if not os.path.exists(path):
                    print(f"[ERROR] Required subdirectory '{path}' is missing.")
                    all_missing = True

        if all_missing:
            print("\n[ERROR] Some required directories are missing. Please ensure all directories are correctly set up.")
        else:
            print("\n[INFO] All required directories are present.")

    print("[INFO] SDXLNude configuration completed.\n")
    return True

def main():
    # Path to the project's root directory (where this script is located)
    project_root = os.path.abspath(os.path.dirname(__file__))

    # Execute FluxCheckpoint configuration
    flux_checkpoint_success = process_flux_checkpoint(project_root)
    if flux_checkpoint_success:
        print("[SUCCESS] FluxCheckpoint executed successfully.\n")
    else:
        print("[FAILURE] FluxCheckpoint encountered errors.\n")

    # Execute FluxLora configuration
    flux_lora_success = process_flux_lora(project_root)
    if flux_lora_success:
        print("[SUCCESS] FluxLora executed successfully.\n")
    else:
        print("[FAILURE] FluxLora encountered errors.\n")

    # Execute SDXLNude configuration
    sdxl_nude_success = process_sdxl_nude(project_root)
    if sdxl_nude_success:
        print("[SUCCESS] SDXLNude executed successfully.\n")
    else:
        print("[FAILURE] SDXLNude encountered errors.\n")

    # Final status
    if flux_checkpoint_success and flux_lora_success and sdxl_nude_success:
        print("[INFO] All configurations have been executed successfully.")
    else:
        print("[INFO] Some configurations encountered errors. Please check the logs above.")

if __name__ == "__main__":
    main()