Update flux_train_ui.py

flux_train_ui.py

@@ -16,6 +16,23 @@ import json
 import yaml
 from slugify import slugify
 from transformers import AutoProcessor, AutoModelForCausalLM
+import logging
+import os
+import yaml
+import uuid
+from slugify import slugify
+import gradio as gr  # Assuming gr is from gradio for error/warning handling
+
+# Configure logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(),  # Output to console
+        logging.FileHandler('training.log')  # Save logs to a file
+    ]
+)
+logger = logging.getLogger(__name__)
 
 sys.path.insert(0, "ai-toolkit")
 from toolkit.job import get_job
@@ -172,47 +189,90 @@ def start_training(
     use_more_advanced_options,
     more_advanced_options,
 ):
+    logger.info("Starting training process")
+    logger.debug(f"Input parameters: lora_name={lora_name}, concept_sentence={concept_sentence}, "
+                 f"steps={steps}, lr={lr}, rank={rank}, model_to_train={model_to_train}, "
+                 f"low_vram={low_vram}, dataset_folder={dataset_folder}, "
+                 f"sample_1={sample_1}, sample_2={sample_2}, sample_3={sample_3}, "
+                 f"use_more_advanced_options={use_more_advanced_options}, "
+                 f"more_advanced_options={more_advanced_options}")
+
     push_to_hub = True
-    print("flux ttain invoke ====================")
+    logger.info("Checking LoRA name")
     if not lora_name:
+        logger.error("LoRA name is empty or None")
         raise gr.Error("You forgot to insert your LoRA name! This name has to be unique.")
+
+    # Check Hugging Face permissions
     try:
-        if whoami()["auth"]["accessToken"]["role"] == "write" or "repo.write" in whoami()["auth"]["accessToken"]["fineGrained"]["scoped"][0]["permissions"]:
-            gr.Info(f"Starting training locally {whoami()['name']}. Your LoRA will be available locally and in Hugging Face after it finishes.")
+        user_info = whoami()
+        logger.debug(f"Hugging Face user info: {user_info}")
+        if user_info["auth"]["accessToken"]["role"] == "write" or \
+           "repo.edit" in user_info["auth"]["accessToken"]["fineGrained"]["scoped"][0]["permissions"]:
+            logger.info(f"Starting training locally for user: {user_info['name']}. LoRA will be available locally and on Hugging Face.")
         else:
             push_to_hub = False
+            logger.warning("No write access to Hugging Face. Training locally only.")
             gr.Warning("Started training locally. Your LoRa will only be available locally because you didn't login with a `write` token to Hugging Face")
-    except:
+    except Exception as e:
         push_to_hub = False
+        logger.error(f"Error checking Hugging Face permissions: {str(e)}")
         gr.Warning("Started training locally. Your LoRa will only be available locally because you didn't login with a `write` token to Hugging Face")
-            
-    print("Started training")
+
+    logger.info("Training started")
     slugged_lora_name = slugify(lora_name)
+    logger.debug(f"Slugged LoRA name: {slugged_lora_name}")
 
     # Load the default config
-    with open("config/examples/train_lora_flux_24gb.yaml", "r") as f:
-        config = yaml.safe_load(f)
+    config_path_default = "config/examples/train_lora_flux_24gb.yaml"
+    logger.info(f"Loading default config from: {config_path_default}")
+    try:
+        with open(config_path_default, "r") as f:
+            config = yaml.safe_load(f)
+        logger.debug(f"Loaded config: {config}")
+    except Exception as e:
+        logger.error(f"Failed to load config from {config_path_default}: {str(e)}")
+        raise
 
     # Update the config with user inputs
-    config["config"]["name"] = slugged_lora_name
-    config["config"]["process"][0]["model"]["low_vram"] = low_vram
-    config["config"]["process"][0]["train"]["skip_first_sample"] = True
-    config["config"]["process"][0]["train"]["steps"] = int(steps)
-    config["config"]["process"][0]["train"]["lr"] = float(lr)
-    config["config"]["process"][0]["network"]["linear"] = int(rank)
-    config["config"]["process"][0]["network"]["linear_alpha"] = int(rank)
-    config["config"]["process"][0]["datasets"][0]["folder_path"] = dataset_folder
-    config["config"]["process"][0]["save"]["push_to_hub"] = push_to_hub
-    if(push_to_hub):
+    logger.info("Updating config with user inputs")
+    try:
+        config["config"]["name"] = slugged_lora_name
+        config["config"]["process"][0]["model"]["low_vram"] = low_vram
+        config["config"]["process"][0]["train"]["skip_first_sample"] = True
+        config["config"]["process"][0]["train"]["steps"] = int(steps)
+        config["config"]["process"][0]["train"]["lr"] = float(lr)
+        config["config"]["process"][0]["network"]["linear"] = int(rank)
+        config["config"]["process"][0]["network"]["linear_alpha"] = int(rank)
+        config["config"]["process"][0]["datasets"][0]["folder_path"] = dataset_folder
+        config["config"]["process"][0]["save"]["push_to_hub"] = push_to_hub
+        logger.debug(f"Updated config fields: name={slugged_lora_name}, low_vram={low_vram}, steps={steps}, "
+                     f"lr={lr}, rank={rank}, dataset_folder={dataset_folder}, push_to_hub={push_to_hub}")
+    except KeyError as e:
+        logger.error(f"Config structure error: Missing key {str(e)}")
+        raise
+    except Exception as e:
+        logger.error(f"Error updating config: {str(e)}")
+        raise
+
+    # Handle Hugging Face repository settings
+    if push_to_hub:
         try:
             username = whoami()["name"]
-        except:
+            logger.debug(f"Hugging Face username: {username}")
+            config["config"]["process"][0]["save"]["hf_repo_id"] = f"{username}/{slugged_lora_name}"
+            config["config"]["process"][0]["save"]["hf_private"] = True
+            logger.debug(f"Set Hugging Face repo: {username}/{slugged_lora_name}")
+        except Exception as e:
+            logger.error(f"Error retrieving Hugging Face username: {str(e)}")
             raise gr.Error("Error trying to retrieve your username. Are you sure you are logged in with Hugging Face?")
-        config["config"]["process"][0]["save"]["hf_repo_id"] = f"{username}/{slugged_lora_name}"
-        config["config"]["process"][0]["save"]["hf_private"] = True
+
+    # Handle concept sentence
     if concept_sentence:
         config["config"]["process"][0]["trigger_word"] = concept_sentence
-    
+        logger.debug(f"Set trigger_word: {concept_sentence}")
+
+    # Handle sampling prompts
     if sample_1 or sample_2 or sample_3:
         config["config"]["process"][0]["train"]["disable_sampling"] = False
         config["config"]["process"][0]["sample"]["sample_every"] = steps
@@ -224,33 +284,56 @@ def start_training(
             config["config"]["process"][0]["sample"]["prompts"].append(sample_2)
         if sample_3:
             config["config"]["process"][0]["sample"]["prompts"].append(sample_3)
+        logger.debug(f"Sampling enabled with prompts: {config['config']['process'][0]['sample']['prompts']}")
     else:
         config["config"]["process"][0]["train"]["disable_sampling"] = True
-    if(model_to_train == "schnell"):
+        logger.debug("Sampling disabled")
+
+    # Handle model selection
+    if model_to_train == "schnell":
         config["config"]["process"][0]["model"]["name_or_path"] = "black-forest-labs/FLUX.1-schnell"
         config["config"]["process"][0]["model"]["assistant_lora_path"] = "ostris/FLUX.1-schnell-training-adapter"
         config["config"]["process"][0]["sample"]["sample_steps"] = 4
-    if(use_more_advanced_options):
-        more_advanced_options_dict = yaml.safe_load(more_advanced_options)
-        config["config"]["process"][0] = recursive_update(config["config"]["process"][0], more_advanced_options_dict)
-        print(config)
-    
+        logger.debug("Using schnell model configuration")
+
+    # Handle advanced options
+    if use_more_advanced_options:
+        try:
+            more_advanced_options_dict = yaml.safe_load(more_advanced_options)
+            logger.debug(f"Advanced options parsed: {more_advanced_options_dict}")
+            config["config"]["process"][0] = recursive_update(config["config"]["process"][0], more_advanced_options_dict)
+            logger.debug(f"Config after advanced options update: {config}")
+        except Exception as e:
+            logger.error(f"Error parsing or applying advanced options: {str(e)}")
+            raise
+
     # Save the updated config
-    # generate a random name for the config
+    logger.info("Saving updated config")
     random_config_name = str(uuid.uuid4())
     os.makedirs("tmp", exist_ok=True)
     config_path = f"tmp/{random_config_name}-{slugged_lora_name}.yaml"
-    with open(config_path, "w") as f:
-        yaml.dump(config, f)
-    
-    
-    
-    print(f"[INFO] Starting training with config: {config_path}")
-    # run the job locally
-    job = get_job(config_path)
-    job.run()
-    job.cleanup()
-
+    try:
+        with open(config_path, "w") as f:
+            yaml.dump(config, f)
+        logger.info(f"Config saved to: {config_path}")
+    except Exception as e:
+        logger.error(f"Error saving config to {config_path}: {str(e)}")
+        raise
+
+    # Run the training job
+    logger.info(f"Starting training job with config: {config_path}")
+    try:
+        job = get_job(config_path)
+        logger.debug("Job object created successfully")
+        job.run()
+        logger.info("Training job completed")
+        job.cleanup()
+        logger.info("Job cleanup completed")
+    except Exception as e:
+        logger.error(f"Error during training job execution: {str(e)}")
+        raise
+
+    logger.info(f"Training completed successfully. Model saved as {slugged_lora_name}")
     return f"Training completed successfully. Model saved as {slugged_lora_name}"
 
 
@@ -267,7 +350,6 @@ def start_training(
 
 
 
-
 config_yaml = '''
 device: cuda:0
 model: