Remove WandB integration from hf_ui.py and run_modal_from_hf.py: eliminated references to WandB logging and related secrets to streamline training process.

hf_ui.py

@@ -115,7 +115,6 @@ def start_training(
     use_more_advanced_options,
     more_advanced_options,
     push_to_hub,
-    use_wandb,
 ):
     print("Starting training from gradio app")
 
@@ -228,7 +227,6 @@ def start_training(
     # add wandb if needed
     config['config']['process'][0]['logging'] = {
         "log_every": 10,
-        "use_wandb": use_wandb,
         "verbose": False
     }
 
@@ -369,7 +367,6 @@ with gr.Blocks(theme=theme, css=css) as demo:
 
         with gr.Row():
             push_to_hub = gr.Checkbox(label="Push to Hub", value=True)
-            use_wandb = gr.Checkbox(label="Use WandB", value=False)
         start = gr.Button("Start training")
         output_components.append(start)
         progress_area = gr.Markdown("")
@@ -408,7 +405,6 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 use_more_advanced_options,
                 more_advanced_options,
                 push_to_hub,
-                use_wandb
             ],
             outputs=progress_area,
         )

run_modal_from_hf.py

@@ -23,7 +23,6 @@ sys.path.insert(0, "/root/ai-toolkit")
 os.environ['DISABLE_TELEMETRY'] = 'YES'
 # Khai báo secrets
 hf_secret = modal.Secret.from_name("huggingface-secret")
-wandb_secret = modal.Secret.from_name("wandb-secret")
 
 # define the volume for storing model outputs, using "creating volumes lazily": https://modal.com/docs/guide/volumes
 # you will find your model, samples and optimizer stored in: https://modal.com/storage/your-username/main/flux-lora-models
@@ -93,7 +92,6 @@ if os.environ.get("DEBUG_TOOLKIT", "0") == "1":
 
 import argparse
 from toolkit.job import get_job
-from toolkit.logging import WandbLogger
 
 def print_end_message(jobs_completed, jobs_failed):
     failure_string = f"{jobs_failed} failure{'' if jobs_failed == 1 else 's'}" if jobs_failed > 0 else ""
@@ -115,7 +113,7 @@ def print_end_message(jobs_completed, jobs_failed):
     gpu="A100", # gpu="H100"
     # more about modal timeouts: https://modal.com/docs/guide/timeouts
     timeout=7200,  # 2 hours, increase or decrease if needed
-    secrets=[hf_secret, wandb_secret]
+    secrets=[hf_secret]
 )
 def main(config_file_list_str: str, recover: bool = False, name: str = None):
     # Các secrets sẽ tự động được inject vào environment variables
@@ -166,21 +164,6 @@ def main(config_file_list_str: str, recover: bool = False, name: str = None):
        os.makedirs(MOUNT_DIR, exist_ok=True)
        print(f"Training outputs will be saved to: {MOUNT_DIR}")
 
-       # setup wandb
-       if config['config']['process'][0]['logging']['use_wandb']:
-           wandb_token = os.environ.get('WANDB_API_KEY', None)
-           if wandb_token:
-               wandb_logger = WandbLogger(
-                   project="flux-lora-training",
-                   run_name=name,
-                   config=job.raw_config,
-               )
-               job.meta["wandb"] = wandb_logger.run.id
-               job.process[0].logger = wandb_logger
-           else:
-               print("WandB token not found, skipping WandB logging")
-               config['config']['process'][0]['logging']['use_wandb'] = False  # disable if no key was given
-       
        # handle dataset zip
        datasets = config['config']['process'][0].get('datasets', None)
        if datasets is not None and isinstance(datasets, list):
@@ -208,9 +191,6 @@ def main(config_file_list_str: str, recover: bool = False, name: str = None):
 
        # run the job
        job.run()
-       
-       if config['config']['process'][0]['logging']['use_wandb']:
-           wandb_logger.finish()
 
        # commit the volume after training
        model_volume.commit()