Update flux_train_ui.py

flux_train_ui.py

@@ -23,17 +23,6 @@ import uuid
 from slugify import slugify
 import gradio as gr  # Assuming gr is from gradio for error/warning handling
 
-os.makedirs("tmp", exist_ok=True)
-# Configure logging
-logging.basicConfig(
-    level=logging.DEBUG,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.StreamHandler(),  # Output to console
-        logging.FileHandler('tmp/training.log')  # Save logs to a file
-    ]
-)
-logger = logging.getLogger(__name__)
 
 sys.path.insert(0, "ai-toolkit")
 from toolkit.job import get_job
@@ -190,8 +179,8 @@ 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}, "
+    print("Starting training process")
+    print(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}, "
@@ -199,44 +188,44 @@ def start_training(
                  f"more_advanced_options={more_advanced_options}")
 
     push_to_hub = True
-    logger.info("Checking LoRA name")
+    print("Checking LoRA name")
     if not lora_name:
-        logger.error("LoRA name is empty or None")
+        print("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:
         user_info = whoami()
-        logger.debug(f"Hugging Face user info: {user_info}")
+        print(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.")
+            print(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.")
+            print("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 Exception as e:
         push_to_hub = False
-        logger.error(f"Error checking Hugging Face permissions: {str(e)}")
+        print(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")
 
-    logger.info("Training started")
+    print("Training started")
     slugged_lora_name = slugify(lora_name)
-    logger.debug(f"Slugged LoRA name: {slugged_lora_name}")
+    print(f"Slugged LoRA name: {slugged_lora_name}")
 
     # Load the default config
     config_path_default = "config/examples/train_lora_flux_24gb.yaml"
-    logger.info(f"Loading default config from: {config_path_default}")
+    print(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}")
+        print(f"Loaded config: {config}")
     except Exception as e:
-        logger.error(f"Failed to load config from {config_path_default}: {str(e)}")
+        print(f"Failed to load config from {config_path_default}: {str(e)}")
         raise
 
     # Update the config with user inputs
-    logger.info("Updating config with user inputs")
+    print("Updating config with user inputs")
     try:
         config["config"]["name"] = slugged_lora_name
         config["config"]["process"][0]["model"]["low_vram"] = low_vram
@@ -247,31 +236,31 @@ def start_training(
         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}, "
+        print(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)}")
+        print(f"Config structure error: Missing key {str(e)}")
         raise
     except Exception as e:
-        logger.error(f"Error updating config: {str(e)}")
+        print(f"Error updating config: {str(e)}")
         raise
 
     # Handle Hugging Face repository settings
     if push_to_hub:
         try:
             username = whoami()["name"]
-            logger.debug(f"Hugging Face username: {username}")
+            print(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}")
+            print(f"Set Hugging Face repo: {username}/{slugged_lora_name}")
         except Exception as e:
-            logger.error(f"Error retrieving Hugging Face username: {str(e)}")
+            print(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?")
 
     # Handle concept sentence
     if concept_sentence:
         config["config"]["process"][0]["trigger_word"] = concept_sentence
-        logger.debug(f"Set trigger_word: {concept_sentence}")
+        print(f"Set trigger_word: {concept_sentence}")
 
     # Handle sampling prompts
     if sample_1 or sample_2 or sample_3:
@@ -285,56 +274,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']}")
+        print(f"Sampling enabled with prompts: {config['config']['process'][0]['sample']['prompts']}")
     else:
         config["config"]["process"][0]["train"]["disable_sampling"] = True
-        logger.debug("Sampling disabled")
+        print("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
-        logger.debug("Using schnell model configuration")
+        print("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}")
+            print(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}")
+            print(f"Config after advanced options update: {config}")
         except Exception as e:
-            logger.error(f"Error parsing or applying advanced options: {str(e)}")
+            print(f"Error parsing or applying advanced options: {str(e)}")
             raise
 
     # Save the updated config
-    logger.info("Saving updated config")
+    print("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"
     try:
         with open(config_path, "w") as f:
             yaml.dump(config, f)
-        logger.info(f"Config saved to: {config_path}")
+        print(f"Config saved to: {config_path}")
     except Exception as e:
-        logger.error(f"Error saving config to {config_path}: {str(e)}")
+        print(f"Error saving config to {config_path}: {str(e)}")
         raise
 
     # Run the training job
-    logger.info(f"Starting training job with config: {config_path}")
+    print(f"Starting training job with config: {config_path}")
     try:
         job = get_job(config_path)
-        logger.debug("Job object created successfully")
+        print("Job object created successfully")
         job.run()
-        logger.info("Training job completed")
+        print("Training job completed")
         job.cleanup()
-        logger.info("Job cleanup completed")
+        print("Job cleanup completed")
     except Exception as e:
-        logger.error(f"Error during training job execution: {str(e)}")
+        print(f"Error during training job execution: {str(e)}")
         raise
 
-    logger.info(f"Training completed successfully. Model saved as {slugged_lora_name}")
+    print(f"Training completed successfully. Model saved as {slugged_lora_name}")
     return f"Training completed successfully. Model saved as {slugged_lora_name}"