Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,36 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_1_step_019799_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_5_step_003300_sample00.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_3_step_000000.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/train_0_step_009900_sample02.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_4_step_019799_sample01.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/train_0_step_019799_sample01.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_3_step_013200_sample01.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_1_step_013200.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/train_0_step_000000_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_0_step_009900_sample04.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/val_1_step_000000_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_5_step_016500_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/train_0_step_019799_sample00.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_3_step_019799_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_0_step_003300.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_3_step_003300_sample02.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/train_1_step_019799.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_2_step_003300_sample04.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/train_0_step_006600_sample01.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_5_step_000000_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/val_1_step_019799_sample04.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_0_step_006600_sample02.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_2_step_006600_sample00.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/val_0_step_013200_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_0_step_009900_sample00.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/train_0_step_009900_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_4_step_000000_sample04.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_2_step_013200_sample04.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/train_1_step_000000_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_2_step_013200_sample00.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/val_1_step_003300_sample03.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/test_3_step_013200_sample00.png filter=lfs diff=lfs merge=lfs -text
+lora_256_dim_512_H100/step_vis/val_0_step_019799_sample04.png filter=lfs diff=lfs merge=lfs -text

lora_baseline_h100/tmux_run.log

@@ -0,0 +1,108 @@
+Run dir : output/lora_baseline_h100
+Log file: output/lora_baseline_h100/train.log
+GPU: NVIDIA GeForce RTX 5090 | VRAM: 31.4 GiB | PyTorch: 2.11.0+cu130
+
+WARNING: WandB init failed — training will continue without logging.
+  Error: No API key configured. Use `wandb login` to log in.
+  Fix: check WANDB_API_KEY permissions, or run with --no-wandb to silence this.
+
+wandb: WARNING Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to https://wandb.me/wandb-init.
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/lora_baseline_h100
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   500
+    batch_size                   4
+    learning_rate                0.0006
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 10
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                lora_baseline_h100
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 11.6/188.5 GiB (6.2%)  |  VRAM: 0.0/31.4 GiB (0.0%)
+
+Pre-flight checks...
+  ✓ torch
+  ✓ diffusers
+  ✓ safetensors
+  ✓ Pillow
+  ✓ tifffile
+  ✓ wandb
+  ✓ transformers
+  ✓ psutil
+  ✓ GPU: NVIDIA GeForce RTX 5090 (31.4 GiB VRAM)
+  ✓ transformer_path: weights/flux2_dev_fp8mixed.safetensors (35.5 GB)
+  ✓ vae_path: weights/flux2-vae.safetensors (0.3 GB)
+  ✓ controlnet_path: weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors (8.2 GB)
+  ✓ text_encoder_path: weights/mistral_3_small_flux2_fp8.safetensors (18.0 GB)
+  ✓ train/rgb: 400 files
+  ✓ train/seg: 400 files
+  ✓ train/depth: 400 files
+  ✓ val/rgb: 80 files
+  ✓ val/seg: 80 files
+  ✓ val/depth: 80 files
+  ✓ test/rgb: 30 files
+  ✓ test/seg: 30 files
+  ✓ test/depth: 30 files
+  ✓ prompt.json found
+
+All pre-flight checks passed.
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Loading cached embedding from output/text_embeddings_global.pt
+  Loaded global text embedding from output/text_embeddings_global.pt (shape: torch.Size([512, 15360]))
+
+============================================================
+  [2/8] Loading VAE
+============================================================
+  Done (1.6s), VRAM: 0.16 GiB
+[MEM @ after VAE] RAM: 11.9/188.5 GiB (6.3%)  |  VRAM: 0.2/31.4 GiB (0.5%)
+
+============================================================
+  [3/8] Loading Transformer
+============================================================
+  Dequantizing FP8 transformer weights...
+Traceback (most recent call last):
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1305, in <module>
+    main()
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 771, in main
+    transformer = load_transformer(
+                  ^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/utility.py", line 143, in load_transformer
+    dequant_sd, fp8_count = dequant_fp8_state_dict(transformer_path, device=device, dtype=dtype)
+                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/utility.py", line 85, in dequant_fp8_state_dict
+    dequant_sd[key] = (tensor.to(torch.float32) * scale).to(dtype).to(device)
+                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 288.00 MiB. GPU 0 has a total capacity of 31.36 GiB of which 102.56 MiB is free. Including non-PyTorch memory, this process has 31.24 GiB memory in use. Of the allocated memory 30.65 GiB is allocated by PyTorch, and 8.62 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)

lora_baseline_h100/train.log

@@ -0,0 +1,1022 @@
+Run dir : output/lora_baseline_h100
+Log file: output/lora_baseline_h100/train.log
+GPU: NVIDIA GeForce RTX 5090 | VRAM: 31.4 GiB | PyTorch: 2.11.0+cu130
+
+WARNING: WandB init failed — training will continue without logging.
+  Error: No API key configured. Use `wandb login` to log in.
+  Fix: check WANDB_API_KEY permissions, or run with --no-wandb to silence this.
+
+wandb: WARNING Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to https://wandb.me/wandb-init.
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/lora_baseline_h100
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   500
+    batch_size                   4
+    learning_rate                0.0006
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 10
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                lora_baseline_h100
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 10.4/188.5 GiB (5.5%)  |  VRAM: 0.0/31.4 GiB (0.0%)
+
+Pre-flight checks...
+  ✓ torch
+  ✓ diffusers
+  ✓ safetensors
+  ✓ Pillow
+  ✓ tifffile
+  ✓ wandb
+  ✓ transformers
+  ✓ psutil
+  ✓ GPU: NVIDIA GeForce RTX 5090 (31.4 GiB VRAM)
+  ✓ transformer_path: weights/flux2_dev_fp8mixed.safetensors (35.5 GB)
+  ✓ vae_path: weights/flux2-vae.safetensors (0.3 GB)
+  ✓ controlnet_path: weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors (8.2 GB)
+  ✓ text_encoder_path: weights/mistral_3_small_flux2_fp8.safetensors (18.0 GB)
+  ✓ train/rgb: 400 files
+  ✓ train/seg: 400 files
+  ✓ train/depth: 400 files
+  ✓ val/rgb: 80 files
+  ✓ val/seg: 80 files
+  ✓ val/depth: 80 files
+  ✓ test/rgb: 30 files
+  ✓ test/seg: 30 files
+  ✓ test/depth: 30 files
+  ✓ prompt.json found
+
+All pre-flight checks passed.
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Loading cached embedding from output/text_embeddings_global.pt
+  Loaded global text embedding from output/text_embeddings_global.pt (shape: torch.Size([512, 15360]))
+
+============================================================
+  [2/8] Loading VAE
+============================================================
+  Done (1.5s), VRAM: 0.16 GiB
+[MEM @ after VAE] RAM: 10.7/188.5 GiB (5.7%)  |  VRAM: 0.2/31.4 GiB (0.5%)
+
+============================================================
+  [3/8] Loading Transformer
+============================================================
+  Dequantizing FP8 transformer weights...
+Traceback (most recent call last):
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1305, in <module>
+    main()
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 771, in main
+    transformer = load_transformer(
+                  ^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/utility.py", line 143, in load_transformer
+    dequant_sd, fp8_count = dequant_fp8_state_dict(transformer_path, device=device, dtype=dtype)
+                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/utility.py", line 85, in dequant_fp8_state_dict
+    dequant_sd[key] = (tensor.to(torch.float32) * scale).to(dtype).to(device)
+                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 288.00 MiB. GPU 0 has a total capacity of 31.36 GiB of which 102.56 MiB is free. Including non-PyTorch memory, this process has 31.24 GiB memory in use. Of the allocated memory 30.65 GiB is allocated by PyTorch, and 8.62 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+Run dir : output/lora_baseline_h100
+Log file: output/lora_baseline_h100/train.log
+GPU: NVIDIA GeForce RTX 5090 | VRAM: 31.4 GiB | PyTorch: 2.11.0+cu130
+
+WARNING: WandB init failed — training will continue without logging.
+  Error: No API key configured. Use `wandb login` to log in.
+  Fix: check WANDB_API_KEY permissions, or run with --no-wandb to silence this.
+
+wandb: WARNING Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to https://wandb.me/wandb-init.
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/lora_baseline_h100
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   500
+    batch_size                   4
+    learning_rate                0.0006
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 10
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                lora_baseline_h100
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 11.6/188.5 GiB (6.2%)  |  VRAM: 0.0/31.4 GiB (0.0%)
+
+Pre-flight checks...
+  ✓ torch
+  ✓ diffusers
+  ✓ safetensors
+  ✓ Pillow
+  ✓ tifffile
+  ✓ wandb
+  ✓ transformers
+  ✓ psutil
+  ✓ GPU: NVIDIA GeForce RTX 5090 (31.4 GiB VRAM)
+  ✓ transformer_path: weights/flux2_dev_fp8mixed.safetensors (35.5 GB)
+  ✓ vae_path: weights/flux2-vae.safetensors (0.3 GB)
+  ✓ controlnet_path: weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors (8.2 GB)
+  ✓ text_encoder_path: weights/mistral_3_small_flux2_fp8.safetensors (18.0 GB)
+  ✓ train/rgb: 400 files
+  ✓ train/seg: 400 files
+  ✓ train/depth: 400 files
+  ✓ val/rgb: 80 files
+  ✓ val/seg: 80 files
+  ✓ val/depth: 80 files
+  ✓ test/rgb: 30 files
+  ✓ test/seg: 30 files
+  ✓ test/depth: 30 files
+  ✓ prompt.json found
+
+All pre-flight checks passed.
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Loading cached embedding from output/text_embeddings_global.pt
+  Loaded global text embedding from output/text_embeddings_global.pt (shape: torch.Size([512, 15360]))
+
+============================================================
+  [2/8] Loading VAE
+============================================================
+  Done (1.6s), VRAM: 0.16 GiB
+[MEM @ after VAE] RAM: 11.9/188.5 GiB (6.3%)  |  VRAM: 0.2/31.4 GiB (0.5%)
+
+============================================================
+  [3/8] Loading Transformer
+============================================================
+  Dequantizing FP8 transformer weights...
+Traceback (most recent call last):
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1305, in <module>
+    main()
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 771, in main
+    transformer = load_transformer(
+                  ^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/utility.py", line 143, in load_transformer
+    dequant_sd, fp8_count = dequant_fp8_state_dict(transformer_path, device=device, dtype=dtype)
+                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/utility.py", line 85, in dequant_fp8_state_dict
+    dequant_sd[key] = (tensor.to(torch.float32) * scale).to(dtype).to(device)
+                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 288.00 MiB. GPU 0 has a total capacity of 31.36 GiB of which 102.56 MiB is free. Including non-PyTorch memory, this process has 31.24 GiB memory in use. Of the allocated memory 30.65 GiB is allocated by PyTorch, and 8.62 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+Run dir : output/lora_baseline_h100
+Log file: output/lora_baseline_h100/train.log
+GPU: NVIDIA H100 NVL | VRAM: 93.1 GiB | PyTorch: 2.11.0+cu130
+wandb: (1) Create a W&B account
+wandb: (2) Use an existing W&B account
+wandb: (3) Don't visualize my results
+wandb: Enter your choice: wandb: You chose 'Create a W&B account'
+wandb: Create an account here: https://wandb.ai/authorize?signup=true&ref=models
+wandb: After creating your account, create a new API key and store it securely.
+wandb: Paste your API key and hit enter:wandb: No netrc file found, creating one.
+wandb: Appending key for api.wandb.ai to your netrc file: /home/xg_wang_group/.netrc
+wandb: Currently logged in as: hkujasonjiang (hku-xg-boost) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Waiting for wandb.init()...
+
+
+
+
+
+
+
+
+
+wandb: Run data is saved locally in /home/xg_wang_group/SynthUrbanSAT/wandb/run-20260416_114715-hu0pyawd
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run lora_baseline_h100-train-20260416-114534
+wandb:  View project at https://wandb.ai/hku-xg-boost/lora_baseline_h100
+wandb:  View run at https://wandb.ai/hku-xg-boost/lora_baseline_h100/runs/hu0pyawd
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/lora_baseline_h100
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   500
+    batch_size                   4
+    learning_rate                0.0006
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 10
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                lora_baseline_h100
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 10.1/188.5 GiB (5.3%)  |  VRAM: 0.0/93.1 GiB (0.0%)
+
+Pre-flight checks...
+  ✓ torch
+  ✓ diffusers
+  ✓ safetensors
+  ✓ Pillow
+  ✓ tifffile
+  ✓ wandb
+  ✓ transformers
+  ✓ psutil
+  ✓ GPU: NVIDIA H100 NVL (93.1 GiB VRAM)
+  ✓ transformer_path: weights/flux2_dev_fp8mixed.safetensors (35.5 GB)
+  ✓ vae_path: weights/flux2-vae.safetensors (0.3 GB)
+  ✓ controlnet_path: weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors (8.2 GB)
+  ✓ text_encoder_path: weights/mistral_3_small_flux2_fp8.safetensors (18.0 GB)
+  ✓ train/rgb: 400 files
+  ✓ train/seg: 400 files
+  ✓ train/depth: 400 files
+  ✓ val/rgb: 80 files
+  ✓ val/seg: 80 files
+  ✓ val/depth: 80 files
+  ✓ test/rgb: 30 files
+  ✓ test/seg: 30 files
+  ✓ test/depth: 30 files
+  ✓ prompt.json found
+
+All pre-flight checks passed.
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Loading cached embedding from output/text_embeddings_global.pt
+  Loaded global text embedding from output/text_embeddings_global.pt (shape: torch.Size([512, 15360]))
+
+============================================================
+  [2/8] Loading VAE
+============================================================
+  Done (1.5s), VRAM: 0.16 GiB
+[MEM @ after VAE] RAM: 10.4/188.5 GiB (5.5%)  |  VRAM: 0.2/93.1 GiB (0.2%)
+
+============================================================
+  [3/8] Loading Transformer
+============================================================
+  Dequantizing FP8 transformer weights...
+  Dequantized 128 FP8 tensors
+  Converting ComfyUI → diffusers keys...
+  Converted: 331 diffusers keys
+  Loading ControlNet weights...
+  ControlNet: 76 keys
+  Creating Flux2ControlTransformer2DModel (control_in_dim=3072)...
+  Skipped 2 control_img_in keys (dim mismatch):
+    control_img_in.bias [6144]
+    control_img_in.weight [6144, 260]
+  Missing: 2, Unexpected: 0
+  Initialized control_img_in.weight [6144, 3072] on cuda
+  Initialized control_img_in.bias [6144] on cuda
+  FP8 compression: 203 frozen Linears, 67.9 → 37.9 GiB (saved 30.0 GiB)
+  Done (57.1s), VRAM: 37.87 GiB
+  Gradient checkpointing: enabled
+  Backbone FROZEN: all transformer params set requires_grad=False
+  Gradients will still propagate to HDC²A via control_context autograd
+[MEM @ after Transformer] RAM: 11.9/188.5 GiB (6.3%)  |  VRAM: 37.9/93.1 GiB (40.7%)
+
+============================================================
+  [4/8] Creating HDC²A Adapter
+============================================================
+HDC²A: 52.4M params
+Control: 0.0M params
+Total trainable: 52.4M params
+
+============================================================
+  [4.5/8] Applying LoRA to ControlNet Control Blocks
+============================================================
+  LoRA rank=32, alpha=32.0, dropout=0
+  LoRA control_transformer_blocks.0.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_out.0 [6144→6144]
+
+  LoRA modules injected: 25
+  LoRA trainable params: 9.83M
+
+Parameter Statistics:
+  HDC²A Adapter:         total=52.4M  trainable=52.4M
+  ControlNet (frozen):   total=4143.4M  LoRA trainable=9.83M
+  Flux2 backbone:        total=0.0M  trainable=0.0M ✓
+  ──────────────────────────────────────────────────
+  Total trainable:       HDC²A 52.4M + LoRA 9.83M = 62.19M
+
+============================================================
+  [5/8] Building Optimizer
+============================================================
+  AdamW: adapter_lr=3.00e-04, backbone_lr=0.00e+00
+    param_group 'adapter': 112 tensors, lr=3.00e-04
+  Scheduler: 400 warmup steps → cosine over ~12500 steps
+  [6/8] Resume: skipped (no checkpoint specified)
+
+============================================================
+  [7/8] Forward Sanity Check
+============================================================
+  [test 1/4] Forward pass (eval mode)...
+    Output shape: torch.Size([1, 4096, 128])
+    Output stats: mean=0.0413, std=0.5156
+    VRAM peak (forward): 68.44 GiB
+  [test 2/4] Loss computation (train mode)...
+    Loss value: 1.335436
+  [test 3/4] Backward pass...
+    Backward completed. VRAM peak (backward): 49.21 GiB
+  [test 4/4] Gradient flow check...
+    HDC²A:   112/112 params have non-zero grad
+    Control: 25/50 params have non-zero grad
+    Top grad norms (HDC²A):
+      semantic_encoder.conv_stem.6.weight: 0.002731
+      depth_encoder.conv_stem.6.weight: 0.002289
+      W_s.weight: 0.002197
+      W_d.weight: 0.002029
+      fusion_blocks.0.ffn_sem.2.weight: 0.001831
+  Test result: PASSED
+[MEM @ after test] RAM: 12.5/188.5 GiB (6.6%)  |  VRAM: 38.1/93.1 GiB (40.9%)
+
+============================================================
+  [8/8] Loading Data
+============================================================
+[Data] Data augmentation: disabled
+[Data] Train: 400 samples, batch_size=4
+[Data] Train: using global text embeddings
+[Data] Val: 80 samples, batch_size=4
+[Data] Test: 30 samples, batch_size=4
+
+======================================================================
+Starting training: 500 epochs × 100 steps = 50000 total steps
+  batch_size=4, grad_accum=4, world_size=1, effective_bs=16
+  adapter_lr=3.00e-04, weight_decay=0.01
+======================================================================
+  [Milestone Vis] steps: [0, 5000, 10000, 15000, 20000, 25000, 29999, 34999, 39999, 44999, 49999]
+  [Milestone Vis] 10 grids: train_0(5), train_1(5), val_0(5), val_1(5), test_0(5), test_1(5), test_2(5), test_3(5), test_4(5), test_5(5)
+
+--- Epoch 0/499  (0% done) ---
+  [MilestoneVis] train_0 step 0 ✓
+  [MilestoneGrid] train_0 → output/lora_baseline_h100/milestone_vis/milestone_grid_train_0.png
+  [MilestoneVis] train_1 step 0 ✓
+  [MilestoneGrid] train_1 → output/lora_baseline_h100/milestone_vis/milestone_grid_train_1.png
+Traceback (most recent call last):
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1305, in <module>
+    main()
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1166, in main
+    train_loss, _epoch_opt_steps = train_one_epoch(
+                                   ^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/train.py", line 250, in train_one_epoch
+    step_vis_callback(global_step, batch)
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1147, in _milestone_step_callback
+    _run_milestone_vis(_global_step)
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1086, in _run_milestone_vis
+    gen_rgb = generate_overfit_samples(
+              ^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
+    return func(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/overfit.py", line 593, in generate_overfit_samples
+    control_context = hdc2a(seg, depth)
+                      ^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/models.py", line 252, in forward
+    T_s, T_d = block(T_s, T_d)
+               ^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/models.py", line 179, in forward
+    k_s = self.rope(k_s)
+          ^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/models.py", line 51, in forward
+    return (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+KeyboardInterrupt
+Exception ignored in atexit callback: <function _start_and_connect_service.<locals>.teardown_atexit at 0x6ffd34c7fce0>
+Traceback (most recent call last):
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/wandb/sdk/lib/service/service_connection.py", line 73, in teardown_atexit
+    conn.teardown(hooks.exit_code)
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/wandb/sdk/lib/service/service_connection.py", line 355, in teardown
+    return self._proc.join()
+           ^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/wandb/sdk/lib/service/service_process.py", line 103, in join
+    return self._process.wait()
+           ^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/subprocess.py", line 1264, in wait
+    return self._wait(timeout=timeout)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/subprocess.py", line 2053, in _wait
+    (pid, sts) = self._try_wait(0)
+                 ^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/subprocess.py", line 2011, in _try_wait
+    (pid, sts) = os.waitpid(self.pid, wait_flags)
+                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+KeyboardInterrupt: 
+Run dir : output/lora_baseline_h100
+Log file: output/lora_baseline_h100/train.log
+GPU: NVIDIA H100 NVL | VRAM: 93.1 GiB | PyTorch: 2.11.0+cu130
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/xg_wang_group/.netrc.
+wandb: Currently logged in as: hkujasonjiang (hku-xg-boost) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Waiting for wandb.init()...
+
+
+
+
+wandb: Run data is saved locally in /home/xg_wang_group/SynthUrbanSAT/wandb/run-20260416_115235-25w3do53
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run lora_baseline_h100-train-20260416-115235
+wandb:  View project at https://wandb.ai/hku-xg-boost/lora_baseline_h100
+wandb:  View run at https://wandb.ai/hku-xg-boost/lora_baseline_h100/runs/25w3do53
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/lora_baseline_h100
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   500
+    batch_size                   4
+    learning_rate                0.0006
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 10
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                lora_baseline_h100
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 10.6/188.5 GiB (5.6%)  |  VRAM: 0.0/93.1 GiB (0.0%)
+
+Pre-flight checks...
+  ✓ torch
+  ✓ diffusers
+  ✓ safetensors
+  ✓ Pillow
+  ✓ tifffile
+  ✓ wandb
+  ✓ transformers
+  ✓ psutil
+  ✓ GPU: NVIDIA H100 NVL (93.1 GiB VRAM)
+  ✓ transformer_path: weights/flux2_dev_fp8mixed.safetensors (35.5 GB)
+  ✓ vae_path: weights/flux2-vae.safetensors (0.3 GB)
+  ✓ controlnet_path: weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors (8.2 GB)
+  ✓ text_encoder_path: weights/mistral_3_small_flux2_fp8.safetensors (18.0 GB)
+  ✓ train/rgb: 400 files
+  ✓ train/seg: 400 files
+  ✓ train/depth: 400 files
+  ✓ val/rgb: 80 files
+  ✓ val/seg: 80 files
+  ✓ val/depth: 80 files
+  ✓ test/rgb: 30 files
+  ✓ test/seg: 30 files
+  ✓ test/depth: 30 files
+  ✓ prompt.json found
+
+All pre-flight checks passed.
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Loading cached embedding from output/text_embeddings_global.pt
+  Loaded global text embedding from output/text_embeddings_global.pt (shape: torch.Size([512, 15360]))
+
+============================================================
+  [2/8] Loading VAE
+============================================================
+  Done (1.5s), VRAM: 0.16 GiB
+[MEM @ after VAE] RAM: 10.9/188.5 GiB (5.8%)  |  VRAM: 0.2/93.1 GiB (0.2%)
+
+============================================================
+  [3/8] Loading Transformer
+============================================================
+  Dequantizing FP8 transformer weights...
+  Dequantized 128 FP8 tensors
+  Converting ComfyUI → diffusers keys...
+  Converted: 331 diffusers keys
+  Loading ControlNet weights...
+  ControlNet: 76 keys
+  Creating Flux2ControlTransformer2DModel (control_in_dim=3072)...
+  Skipped 2 control_img_in keys (dim mismatch):
+    control_img_in.bias [6144]
+    control_img_in.weight [6144, 260]
+  Missing: 2, Unexpected: 0
+  Initialized control_img_in.weight [6144, 3072] on cuda
+  Initialized control_img_in.bias [6144] on cuda
+  FP8 compression: 203 frozen Linears, 67.9 → 37.9 GiB (saved 30.0 GiB)
+  Done (54.9s), VRAM: 37.87 GiB
+  Gradient checkpointing: enabled
+  Backbone FROZEN: all transformer params set requires_grad=False
+  Gradients will still propagate to HDC²A via control_context autograd
+[MEM @ after Transformer] RAM: 12.2/188.5 GiB (6.5%)  |  VRAM: 37.9/93.1 GiB (40.7%)
+
+============================================================
+  [4/8] Creating HDC²A Adapter
+============================================================
+HDC²A: 52.4M params
+Control: 0.0M params
+Total trainable: 52.4M params
+
+============================================================
+  [4.5/8] Applying LoRA to ControlNet Control Blocks
+============================================================
+  LoRA rank=32, alpha=32.0, dropout=0
+  LoRA control_transformer_blocks.0.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_out.0 [6144→6144]
+
+  LoRA modules injected: 25
+  LoRA trainable params: 9.83M
+
+Parameter Statistics:
+  HDC²A Adapter:         total=52.4M  trainable=52.4M
+  ControlNet (frozen):   total=4143.4M  LoRA trainable=9.83M
+  Flux2 backbone:        total=0.0M  trainable=0.0M ✓
+  ──────────────────────────────────────────────────
+  Total trainable:       HDC²A 52.4M + LoRA 9.83M = 62.19M
+
+============================================================
+  [5/8] Building Optimizer
+============================================================
+  AdamW: adapter_lr=3.00e-04, backbone_lr=0.00e+00
+    param_group 'adapter': 112 tensors, lr=3.00e-04
+  Scheduler: 400 warmup steps → cosine over ~12500 steps
+  [6/8] Resume: skipped (no checkpoint specified)
+
+============================================================
+  [7/8] Forward Sanity Check
+============================================================
+  [test 1/4] Forward pass (eval mode)...
+    Output shape: torch.Size([1, 4096, 128])
+    Output stats: mean=0.0413, std=0.5156
+    VRAM peak (forward): 68.44 GiB
+  [test 2/4] Loss computation (train mode)...
+    Loss value: 1.335436
+  [test 3/4] Backward pass...
+    Backward completed. VRAM peak (backward): 49.21 GiB
+  [test 4/4] Gradient flow check...
+    HDC²A:   112/112 params have non-zero grad
+    Control: 25/50 params have non-zero grad
+    Top grad norms (HDC²A):
+      semantic_encoder.conv_stem.6.weight: 0.002731
+      depth_encoder.conv_stem.6.weight: 0.002289
+      W_s.weight: 0.002182
+      W_d.weight: 0.002029
+      fusion_blocks.0.ffn_sem.2.weight: 0.001831
+  Test result: PASSED
+[MEM @ after test] RAM: 12.7/188.5 GiB (6.8%)  |  VRAM: 38.1/93.1 GiB (40.9%)
+
+============================================================
+  [8/8] Loading Data
+============================================================
+[Data] Data augmentation: disabled
+[Data] Train: 400 samples, batch_size=4
+[Data] Train: using global text embeddings
+[Data] Val: 80 samples, batch_size=4
+[Data] Test: 30 samples, batch_size=4
+
+======================================================================
+Starting training: 500 epochs × 100 steps = 50000 total steps
+  batch_size=4, grad_accum=4, world_size=1, effective_bs=16
+  adapter_lr=3.00e-04, weight_decay=0.01
+======================================================================
+  [Milestone Vis] steps: [0, 5000, 10000, 15000, 20000, 25000, 29999, 34999, 39999, 44999, 49999]
+  [Milestone Vis] 10 grids: train_0(5), train_1(5), val_0(5), val_1(5), test_0(5), test_1(5), test_2(5), test_3(5), test_4(5), test_5(5)
+
+--- Epoch 0/499  (0% done) ---
+  [MilestoneVis] train_0 step 0 ✓
+  [MilestoneGrid] train_0 → output/lora_baseline_h100/milestone_vis/milestone_grid_train_0.png
+  [MilestoneVis] train_1 step 0 ✓
+  [MilestoneGrid] train_1 → output/lora_baseline_h100/milestone_vis/milestone_grid_train_1.png
+  [MilestoneVis] val_0 step 0 ✓
+  [MilestoneGrid] val_0 → output/lora_baseline_h100/milestone_vis/milestone_grid_val_0.png
+  [MilestoneVis] val_1 step 0 ✓
+  [MilestoneGrid] val_1 → output/lora_baseline_h100/milestone_vis/milestone_grid_val_1.png
+  [MilestoneVis] test_0 step 0 ✓
+  [MilestoneGrid] test_0 → output/lora_baseline_h100/milestone_vis/milestone_grid_test_0.png
+  [MilestoneVis] test_1 step 0 ✓
+  [MilestoneGrid] test_1 → output/lora_baseline_h100/milestone_vis/milestone_grid_test_1.png
+  [MilestoneVis] test_2 step 0 ✓
+  [MilestoneGrid] test_2 → output/lora_baseline_h100/milestone_vis/milestone_grid_test_2.png
+  [MilestoneVis] test_3 step 0 ✓
+  [MilestoneGrid] test_3 → output/lora_baseline_h100/milestone_vis/milestone_grid_test_3.png
+  [MilestoneVis] test_4 step 0 ✓
+  [MilestoneGrid] test_4 → output/lora_baseline_h100/milestone_vis/milestone_grid_test_4.png
+  [MilestoneVis] test_5 step 0 ✓
+  [MilestoneGrid] test_5 → output/lora_baseline_h100/milestone_vis/milestone_grid_test_5.png
+  [HF Upload WARN] upload failed: No module named 'ControlNet_training'
+  [Epoch 0][10/100] loss=0.941065 avg=0.924597 VRAM=38.6GiB | 0.0% done | ETA(epoch): 10631s
+  [Epoch 0][20/100] loss=0.800769 avg=0.912605 VRAM=38.5GiB | 0.0% done | ETA(epoch): 5255s
+  [Epoch 0][30/100] loss=0.999155 avg=0.911449 VRAM=38.6GiB | 0.1% done | ETA(epoch): 3384s
+  [Epoch 0][40/100] loss=0.830989 avg=0.906137 VRAM=38.5GiB | 0.1% done | ETA(epoch): 2374s
+  [Epoch 0][50/100] loss=0.775814 avg=0.907010 VRAM=38.6GiB | 0.1% done | ETA(epoch): 1716s
+  [Epoch 0][60/100] loss=0.878590 avg=0.908445 VRAM=38.5GiB | 0.1% done | ETA(epoch): 1232s
+  [Epoch 0][70/100] loss=1.031628 avg=0.911537 VRAM=38.6GiB | 0.1% done | ETA(epoch): 849s
+  [Epoch 0][80/100] loss=0.866052 avg=0.914041 VRAM=38.5GiB | 0.2% done | ETA(epoch): 528s
+  [Epoch 0][90/100] loss=0.915645 avg=0.913495 VRAM=38.6GiB | 0.2% done | ETA(epoch): 250s
+  [Epoch 0][100/100] loss=0.869309 avg=0.913075 VRAM=38.5GiB | 0.2% done | ETA(epoch): 0s
+  Train loss: 0.913075 (2379.3s)  ETA: 19824min
+  Val loss: 0.938541  [t_0.0-0.2=1.1016  t_0.2-0.4=1.0422  t_0.4-0.6=0.8825  t_0.6-0.8=0.7715  t_0.8-1.0=0.9322]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0000  (BEST)
+[MEM @ epoch 0 end] RAM: 15.4/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 1/499  (0% done) ---
+  [Epoch 1][10/100] loss=0.809276 avg=0.905533 VRAM=38.6GiB | 0.2% done | ETA(epoch): 1193s
+  [Epoch 1][20/100] loss=0.830221 avg=0.909845 VRAM=38.5GiB | 0.2% done | ETA(epoch): 1062s
+  [Epoch 1][30/100] loss=0.944313 avg=0.909747 VRAM=38.6GiB | 0.3% done | ETA(epoch): 929s
+  [Epoch 1][40/100] loss=0.910843 avg=0.912359 VRAM=38.5GiB | 0.3% done | ETA(epoch): 796s
+  [Epoch 1][50/100] loss=0.963525 avg=0.911716 VRAM=38.6GiB | 0.3% done | ETA(epoch): 664s
+  [Epoch 1][60/100] loss=0.984976 avg=0.909685 VRAM=38.5GiB | 0.3% done | ETA(epoch): 531s
+  [Epoch 1][70/100] loss=0.803163 avg=0.910757 VRAM=38.6GiB | 0.3% done | ETA(epoch): 398s
+  [Epoch 1][80/100] loss=0.954126 avg=0.915441 VRAM=38.5GiB | 0.4% done | ETA(epoch): 266s
+  [Epoch 1][90/100] loss=0.985108 avg=0.914229 VRAM=38.6GiB | 0.4% done | ETA(epoch): 133s
+  [Epoch 1][100/100] loss=1.071944 avg=0.915659 VRAM=38.5GiB | 0.4% done | ETA(epoch): 0s
+  Train loss: 0.915659 (1327.6s)  ETA: 15810min
+  Val loss: 0.930462  [t_0.0-0.2=1.0887  t_0.2-0.4=1.0227  t_0.4-0.6=0.8967  t_0.6-0.8=0.7742  t_0.8-1.0=0.9710]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0001  (BEST)
+[MEM @ epoch 1 end] RAM: 15.5/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 2/499  (0% done) ---
+  [Epoch 2][10/100] loss=0.896954 avg=0.922830 VRAM=38.6GiB | 0.4% done | ETA(epoch): 1194s
+  [Epoch 2][20/100] loss=0.879038 avg=0.923363 VRAM=38.5GiB | 0.4% done | ETA(epoch): 1061s
+  [Epoch 2][30/100] loss=0.824969 avg=0.915989 VRAM=38.6GiB | 0.5% done | ETA(epoch): 929s
+  [Epoch 2][40/100] loss=0.882990 avg=0.917364 VRAM=38.5GiB | 0.5% done | ETA(epoch): 796s
+  [Epoch 2][50/100] loss=1.054123 avg=0.913184 VRAM=38.6GiB | 0.5% done | ETA(epoch): 663s
+  [Epoch 2][60/100] loss=0.863263 avg=0.912003 VRAM=38.5GiB | 0.5% done | ETA(epoch): 531s
+  [Epoch 2][70/100] loss=0.982849 avg=0.912602 VRAM=38.6GiB | 0.5% done | ETA(epoch): 398s
+  [Epoch 2][80/100] loss=0.842028 avg=0.913022 VRAM=38.5GiB | 0.6% done | ETA(epoch): 265s
+  [Epoch 2][90/100] loss=0.888870 avg=0.917269 VRAM=38.6GiB | 0.6% done | ETA(epoch): 133s
+  [Epoch 2][100/100] loss=0.917875 avg=0.915649 VRAM=38.5GiB | 0.6% done | ETA(epoch): 0s
+  Train loss: 0.915649 (1330.5s)  ETA: 14464min
+  Val loss: 0.932581  [t_0.0-0.2=1.1004  t_0.2-0.4=1.0259  t_0.4-0.6=0.8721  t_0.6-0.8=0.7633  t_0.8-1.0=0.9520]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0002  
+[MEM @ epoch 2 end] RAM: 15.5/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 3/499  (1% done) ---
+  [Epoch 3][10/100] loss=0.798546 avg=0.892938 VRAM=38.6GiB | 0.6% done | ETA(epoch): 1192s
+  [Epoch 3][20/100] loss=0.986750 avg=0.928209 VRAM=38.5GiB | 0.6% done | ETA(epoch): 1061s
+  [Epoch 3][30/100] loss=0.774730 avg=0.905527 VRAM=38.6GiB | 0.7% done | ETA(epoch): 928s
+  [Epoch 3][40/100] loss=1.007182 avg=0.905173 VRAM=38.5GiB | 0.7% done | ETA(epoch): 796s
+  [Epoch 3][50/100] loss=0.976228 avg=0.912917 VRAM=38.6GiB | 0.7% done | ETA(epoch): 663s
+  [Epoch 3][60/100] loss=0.981444 avg=0.913284 VRAM=38.5GiB | 0.7% done | ETA(epoch): 530s
+  [Epoch 3][70/100] loss=0.820385 avg=0.909043 VRAM=38.6GiB | 0.7% done | ETA(epoch): 398s
+  [Epoch 3][80/100] loss=0.814420 avg=0.905657 VRAM=38.5GiB | 0.8% done | ETA(epoch): 265s
+  [Epoch 3][90/100] loss=0.834756 avg=0.908909 VRAM=38.6GiB | 0.8% done | ETA(epoch): 133s
+  [Epoch 3][100/100] loss=1.017020 avg=0.909447 VRAM=38.5GiB | 0.8% done | ETA(epoch): 0s
+  Train loss: 0.909447 (1326.2s)  ETA: 13770min
+  Val loss: 0.954529  [t_0.0-0.2=1.0967  t_0.2-0.4=1.0011  t_0.4-0.6=0.8903  t_0.6-0.8=0.7404  t_0.8-1.0=0.9429]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0003  
+  Deleted old checkpoint: checkpoint_epoch_0000
+[MEM @ epoch 3 end] RAM: 15.4/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 4/499  (1% done) ---
+  [Epoch 4][10/100] loss=0.863603 avg=0.873102 VRAM=38.6GiB | 0.8% done | ETA(epoch): 1193s
+  [Epoch 4][20/100] loss=0.797668 avg=0.897814 VRAM=38.5GiB | 0.8% done | ETA(epoch): 1061s
+  [Epoch 4][30/100] loss=0.956929 avg=0.897195 VRAM=38.6GiB | 0.9% done | ETA(epoch): 929s
+  [Epoch 4][40/100] loss=0.987086 avg=0.901464 VRAM=38.5GiB | 0.9% done | ETA(epoch): 796s
+  [Epoch 4][50/100] loss=0.827385 avg=0.904161 VRAM=38.6GiB | 0.9% done | ETA(epoch): 668s
+  [Epoch 4][60/100] loss=0.785728 avg=0.902909 VRAM=38.5GiB | 0.9% done | ETA(epoch): 534s
+  [Epoch 4][70/100] loss=0.924116 avg=0.904073 VRAM=38.6GiB | 0.9% done | ETA(epoch): 400s
+  [Epoch 4][80/100] loss=0.803847 avg=0.898653 VRAM=38.5GiB | 1.0% done | ETA(epoch): 267s
+  [Epoch 4][90/100] loss=0.906517 avg=0.897017 VRAM=38.6GiB | 1.0% done | ETA(epoch): 133s
+  [Epoch 4][100/100] loss=0.838926 avg=0.895190 VRAM=38.5GiB | 1.0% done | ETA(epoch): 0s
+  Train loss: 0.895190 (1332.2s)  ETA: 13354min
+  Val loss: 0.938578  [t_0.0-0.2=1.0855  t_0.2-0.4=1.0255  t_0.4-0.6=0.8686  t_0.6-0.8=0.7556  t_0.8-1.0=0.9351]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0004  
+[MEM @ epoch 4 end] RAM: 15.5/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 5/499  (1% done) ---
+  [Epoch 5][10/100] loss=0.880699 avg=0.898051 VRAM=38.6GiB | 1.0% done | ETA(epoch): 1192s
+  [Epoch 5][20/100] loss=1.052332 avg=0.889742 VRAM=38.5GiB | 1.0% done | ETA(epoch): 1060s
+  [Epoch 5][30/100] loss=0.902173 avg=0.890994 VRAM=38.6GiB | 1.1% done | ETA(epoch): 928s
+  [Epoch 5][40/100] loss=0.933784 avg=0.893495 VRAM=38.5GiB | 1.1% done | ETA(epoch): 801s
+  [Epoch 5][50/100] loss=0.851236 avg=0.892024 VRAM=38.6GiB | 1.1% done | ETA(epoch): 667s
+  [Epoch 5][60/100] loss=0.827588 avg=0.896425 VRAM=38.5GiB | 1.1% done | ETA(epoch): 533s
+  [Epoch 5][70/100] loss=0.910750 avg=0.897489 VRAM=38.6GiB | 1.1% done | ETA(epoch): 399s
+  [Epoch 5][80/100] loss=0.850334 avg=0.894766 VRAM=38.5GiB | 1.2% done | ETA(epoch): 266s
+  [Epoch 5][90/100] loss=0.902218 avg=0.895331 VRAM=38.6GiB | 1.2% done | ETA(epoch): 133s
+  [Epoch 5][100/100] loss=0.984387 avg=0.894165 VRAM=38.5GiB | 1.2% done | ETA(epoch): 0s
+  Train loss: 0.894165 (1330.1s)  ETA: 13066min
+  Val loss: 0.915973  [t_0.0-0.2=1.0750  t_0.2-0.4=1.0342  t_0.4-0.6=0.8432  t_0.6-0.8=0.7724  t_0.8-1.0=0.9757]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0005  (BEST)
+  Deleted old checkpoint: checkpoint_epoch_0001
+  Deleted old checkpoint: checkpoint_epoch_0002
+[MEM @ epoch 5 end] RAM: 15.3/188.5 GiB (8.1%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 6/499  (1% done) ---
+  [Epoch 6][10/100] loss=0.998760 avg=0.877657 VRAM=38.6GiB | 1.2% done | ETA(epoch): 1194s
+  [Epoch 6][20/100] loss=0.828451 avg=0.886384 VRAM=38.5GiB | 1.2% done | ETA(epoch): 1063s
+  [Epoch 6][30/100] loss=0.846599 avg=0.895307 VRAM=38.6GiB | 1.3% done | ETA(epoch): 930s
+  [Epoch 6][40/100] loss=0.897369 avg=0.900879 VRAM=38.5GiB | 1.3% done | ETA(epoch): 797s
+  [Epoch 6][50/100] loss=0.937362 avg=0.898990 VRAM=38.6GiB | 1.3% done | ETA(epoch): 664s
+  [Epoch 6][60/100] loss=0.871279 avg=0.899144 VRAM=38.5GiB | 1.3% done | ETA(epoch): 531s
+  [Epoch 6][70/100] loss=0.882198 avg=0.899551 VRAM=38.6GiB | 1.3% done | ETA(epoch): 398s
+  [Epoch 6][80/100] loss=0.924926 avg=0.897149 VRAM=38.5GiB | 1.4% done | ETA(epoch): 265s
+  [Epoch 6][90/100] loss=0.879220 avg=0.898957 VRAM=38.6GiB | 1.4% done | ETA(epoch): 133s
+  [Epoch 6][100/100] loss=0.926145 avg=0.902229 VRAM=38.5GiB | 1.4% done | ETA(epoch): 0s
+  Train loss: 0.902229 (1327.3s)  ETA: 12851min
+  Val loss: 0.919442  [t_0.0-0.2=1.1017  t_0.2-0.4=1.0331  t_0.4-0.6=0.8386  t_0.6-0.8=0.7633  t_0.8-1.0=0.8907]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0006  
+  Deleted old checkpoint: checkpoint_epoch_0003
+[MEM @ epoch 6 end] RAM: 15.6/188.5 GiB (8.3%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 7/499  (1% done) ---
+  [Epoch 7][10/100] loss=0.861784 avg=0.881602 VRAM=38.6GiB | 1.4% done | ETA(epoch): 1192s
+  [Epoch 7][20/100] loss=0.919683 avg=0.878720 VRAM=38.5GiB | 1.4% done | ETA(epoch): 1061s
+  [Epoch 7][30/100] loss=0.962979 avg=0.882661 VRAM=38.6GiB | 1.5% done | ETA(epoch): 928s
+  [Epoch 7][40/100] loss=0.930568 avg=0.884640 VRAM=38.5GiB | 1.5% done | ETA(epoch): 796s
+  [Epoch 7][50/100] loss=0.974388 avg=0.884026 VRAM=38.6GiB | 1.5% done | ETA(epoch): 663s
+  [Epoch 7][60/100] loss=0.899950 avg=0.882088 VRAM=38.5GiB | 1.5% done | ETA(epoch): 530s
+  [Epoch 7][70/100] loss=0.949109 avg=0.885325 VRAM=38.6GiB | 1.5% done | ETA(epoch): 398s
+  [Epoch 7][80/100] loss=0.953958 avg=0.885340 VRAM=38.5GiB | 1.6% done | ETA(epoch): 265s
+  [Epoch 7][90/100] loss=0.912331 avg=0.888319 VRAM=38.6GiB | 1.6% done | ETA(epoch): 133s
+  [Epoch 7][100/100] loss=0.892251 avg=0.891891 VRAM=38.5GiB | 1.6% done | ETA(epoch): 0s
+  Train loss: 0.891891 (1326.3s)  ETA: 12683min
+  Val loss: 0.948828  [t_0.0-0.2=1.0931  t_0.2-0.4=1.0054  t_0.4-0.6=0.8854  t_0.6-0.8=0.7767  t_0.8-1.0=0.9222]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0007  
+  Deleted old checkpoint: checkpoint_epoch_0004
+[MEM @ epoch 7 end] RAM: 15.4/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 8/499  (2% done) ---
+  [Epoch 8][10/100] loss=0.793007 avg=0.882597 VRAM=38.6GiB | 1.6% done | ETA(epoch): 1194s
+  [Epoch 8][20/100] loss=0.778859 avg=0.884551 VRAM=38.5GiB | 1.6% done | ETA(epoch): 1061s
+  [Epoch 8][30/100] loss=0.949995 avg=0.887610 VRAM=38.6GiB | 1.7% done | ETA(epoch): 928s
+  [Epoch 8][40/100] loss=0.959300 avg=0.894274 VRAM=38.5GiB | 1.7% done | ETA(epoch): 796s
+  [Epoch 8][50/100] loss=0.950433 avg=0.897313 VRAM=38.6GiB | 1.7% done | ETA(epoch): 663s
+  [Epoch 8][60/100] loss=0.888023 avg=0.898194 VRAM=38.5GiB | 1.7% done | ETA(epoch): 531s
+  [Epoch 8][70/100] loss=0.898501 avg=0.898381 VRAM=38.6GiB | 1.7% done | ETA(epoch): 398s
+  [Epoch 8][80/100] loss=0.839228 avg=0.900195 VRAM=38.5GiB | 1.8% done | ETA(epoch): 266s
+  [Epoch 8][90/100] loss=0.945698 avg=0.898329 VRAM=38.6GiB | 1.8% done | ETA(epoch): 133s
+  [Epoch 8][100/100] loss=0.846462 avg=0.896364 VRAM=38.5GiB | 1.8% done | ETA(epoch): 0s
+  Train loss: 0.896364 (1328.6s)  ETA: 12549min
+  Val loss: 0.929101  [t_0.0-0.2=1.0751  t_0.2-0.4=1.0113  t_0.4-0.6=0.8559  t_0.6-0.8=0.7663  t_0.8-1.0=0.9477]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0008  
+[MEM @ epoch 8 end] RAM: 15.5/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 9/499  (2% done) ---
+  [Epoch 9][10/100] loss=0.811317 avg=0.850544 VRAM=38.6GiB | 1.8% done | ETA(epoch): 1192s
+  [Epoch 9][20/100] loss=0.878190 avg=0.864503 VRAM=38.5GiB | 1.8% done | ETA(epoch): 1060s
+  [Epoch 9][30/100] loss=0.816290 avg=0.875198 VRAM=38.6GiB | 1.9% done | ETA(epoch): 928s
+  [Epoch 9][40/100] loss=0.899423 avg=0.874578 VRAM=38.5GiB | 1.9% done | ETA(epoch): 795s
+  [Epoch 9][50/100] loss=0.765660 avg=0.879674 VRAM=38.6GiB | 1.9% done | ETA(epoch): 663s
+  [Epoch 9][60/100] loss=1.041576 avg=0.888781 VRAM=38.5GiB | 1.9% done | ETA(epoch): 534s
+  [Epoch 9][70/100] loss=0.967851 avg=0.896216 VRAM=38.6GiB | 1.9% done | ETA(epoch): 400s
+  [Epoch 9][80/100] loss=0.869709 avg=0.898694 VRAM=38.5GiB | 2.0% done | ETA(epoch): 267s
+  [Epoch 9][90/100] loss=0.913030 avg=0.894023 VRAM=38.6GiB | 2.0% done | ETA(epoch): 133s
+  [Epoch 9][100/100] loss=0.905419 avg=0.890292 VRAM=38.5GiB | 2.0% done | ETA(epoch): 0s
+  Train loss: 0.890292 (1332.6s)  ETA: 12440min
+  Val loss: 0.907056  [t_0.0-0.2=1.0759  t_0.2-0.4=1.0022  t_0.4-0.6=0.8641  t_0.6-0.8=0.7658  t_0.8-1.0=0.8648]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0009  (BEST)
+  Deleted old checkpoint: checkpoint_epoch_0005
+  Deleted old checkpoint: checkpoint_epoch_0006
+[MEM @ epoch 9 end] RAM: 15.4/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 10/499  (2% done) ---
+  [Epoch 10][10/100] loss=0.799796 avg=0.875355 VRAM=38.6GiB | 2.0% done | ETA(epoch): 1193s
+  [Epoch 10][20/100] loss=0.885095 avg=0.887158 VRAM=38.5GiB | 2.0% done | ETA(epoch): 1061s
+  [Epoch 10][30/100] loss=0.908264 avg=0.882971 VRAM=38.6GiB | 2.1% done | ETA(epoch): 928s
+  [Epoch 10][40/100] loss=0.998384 avg=0.885364 VRAM=38.5GiB | 2.1% done | ETA(epoch): 795s
+  [Epoch 10][50/100] loss=0.967074 avg=0.892462 VRAM=38.6GiB | 2.1% done | ETA(epoch): 665s
+  [Epoch 10][60/100] loss=0.807719 avg=0.891286 VRAM=38.5GiB | 2.1% done | ETA(epoch): 532s
+  [Epoch 10][70/100] loss=0.930732 avg=0.894307 VRAM=38.6GiB | 2.1% done | ETA(epoch): 399s
+  [Epoch 10][80/100] loss=0.934865 avg=0.891560 VRAM=38.5GiB | 2.2% done | ETA(epoch): 266s
+  [Epoch 10][90/100] loss=0.879800 avg=0.896282 VRAM=38.6GiB | 2.2% done | ETA(epoch): 133s
+  [Epoch 10][100/100] loss=0.951419 avg=0.893173 VRAM=38.5GiB | 2.2% done | ETA(epoch): 0s
+  Train loss: 0.893173 (1328.5s)  ETA: 12343min
+  Val loss: 0.915998  [t_0.0-0.2=1.0966  t_0.2-0.4=1.0120  t_0.4-0.6=0.8423  t_0.6-0.8=0.7475  t_0.8-1.0=0.9145]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0010  
+  Deleted old checkpoint: checkpoint_epoch_0007
+[MEM @ epoch 10 end] RAM: 15.4/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 11/499  (2% done) ---
+  [Epoch 11][10/100] loss=0.898770 avg=0.840691 VRAM=38.6GiB | 2.2% done | ETA(epoch): 1192s
+  [Epoch 11][20/100] loss=0.902107 avg=0.868226 VRAM=38.5GiB | 2.2% done | ETA(epoch): 1061s
+  [Epoch 11][30/100] loss=0.899334 avg=0.886091 VRAM=38.6GiB | 2.3% done | ETA(epoch): 929s
+  [Epoch 11][40/100] loss=0.836345 avg=0.883065 VRAM=38.5GiB | 2.3% done | ETA(epoch): 796s
+  [Epoch 11][50/100] loss=0.838135 avg=0.884279 VRAM=38.6GiB | 2.3% done | ETA(epoch): 663s
+  [Epoch 11][60/100] loss=0.886633 avg=0.885590 VRAM=38.5GiB | 2.3% done | ETA(epoch): 531s
+  [Epoch 11][70/100] loss=0.935280 avg=0.888742 VRAM=38.6GiB | 2.3% done | ETA(epoch): 398s
+  [Epoch 11][80/100] loss=0.892184 avg=0.887628 VRAM=38.5GiB | 2.4% done | ETA(epoch): 265s
+  [Epoch 11][90/100] loss=0.885992 avg=0.887280 VRAM=38.6GiB | 2.4% done | ETA(epoch): 133s
+  [Epoch 11][100/100] loss=0.856227 avg=0.888817 VRAM=38.5GiB | 2.4% done | ETA(epoch): 0s
+  Train loss: 0.888817 (1326.6s)  ETA: 12257min
+  Val loss: 0.924249  [t_0.0-0.2=1.0761  t_0.2-0.4=1.0630  t_0.4-0.6=0.8336  t_0.6-0.8=0.7874  t_0.8-1.0=0.8219]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0011  
+  Deleted old checkpoint: checkpoint_epoch_0008
+[MEM @ epoch 11 end] RAM: 15.5/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 12/499  (2% done) ---
+  [Epoch 12][10/100] loss=0.879421 avg=0.851119 VRAM=38.6GiB | 2.4% done | ETA(epoch): 1193s
+  [Epoch 12][20/100] loss=0.906359 avg=0.879505 VRAM=38.5GiB | 2.4% done | ETA(epoch): 1061s
+  [Epoch 12][30/100] loss=0.807515 avg=0.879843 VRAM=38.6GiB | 2.5% done | ETA(epoch): 929s
+  [Epoch 12][40/100] loss=0.951137 avg=0.893757 VRAM=38.5GiB | 2.5% done | ETA(epoch): 796s
+  [Epoch 12][50/100] loss=1.001705 avg=0.896496 VRAM=38.6GiB | 2.5% done | ETA(epoch): 663s
+  [Epoch 12][60/100] loss=0.866891 avg=0.892637 VRAM=38.5GiB | 2.5% done | ETA(epoch): 530s
+  [Epoch 12][70/100] loss=0.944366 avg=0.892345 VRAM=38.6GiB | 2.5% done | ETA(epoch): 398s
+  [Epoch 12][80/100] loss=0.969197 avg=0.891150 VRAM=38.5GiB | 2.6% done | ETA(epoch): 265s
+  [Epoch 12][90/100] loss=0.948820 avg=0.894898 VRAM=38.6GiB | 2.6% done | ETA(epoch): 133s
+  [Epoch 12][100/100] loss=0.967688 avg=0.894570 VRAM=38.5GiB | 2.6% done | ETA(epoch): 0s
+  Train loss: 0.894570 (1325.8s)  ETA: 12181min
+  Val loss: 0.910472  [t_0.0-0.2=1.0759  t_0.2-0.4=0.9960  t_0.4-0.6=0.8730  t_0.6-0.8=0.7371  t_0.8-1.0=0.8862]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0012  
+[MEM @ epoch 12 end] RAM: 15.5/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 13/499  (3% done) ---
+  [Epoch 13][10/100] loss=0.870758 avg=0.880895 VRAM=38.6GiB | 2.6% done | ETA(epoch): 1193s
+  [Epoch 13][20/100] loss=0.939617 avg=0.881841 VRAM=38.5GiB | 2.6% done | ETA(epoch): 1061s
+  [Epoch 13][30/100] loss=0.893678 avg=0.891095 VRAM=38.6GiB | 2.7% done | ETA(epoch): 929s
+  [Epoch 13][40/100] loss=0.894570 avg=0.884680 VRAM=38.5GiB | 2.7% done | ETA(epoch): 796s
+  [Epoch 13][50/100] loss=0.818397 avg=0.888858 VRAM=38.6GiB | 2.7% done | ETA(epoch): 663s
+  [Epoch 13][60/100] loss=1.002739 avg=0.890945 VRAM=38.5GiB | 2.7% done | ETA(epoch): 531s
+  [Epoch 13][70/100] loss=0.937253 avg=0.890811 VRAM=38.6GiB | 2.7% done | ETA(epoch): 398s
+  [Epoch 13][80/100] loss=0.975271 avg=0.894329 VRAM=38.5GiB | 2.8% done | ETA(epoch): 265s
+  [Epoch 13][90/100] loss=0.830887 avg=0.896053 VRAM=38.6GiB | 2.8% done | ETA(epoch): 133s
+  [Epoch 13][100/100] loss=0.878928 avg=0.899050 VRAM=38.5GiB | 2.8% done | ETA(epoch): 0s
+  Train loss: 0.899050 (1329.5s)  ETA: 12113min
+  Val loss: 0.918244  [t_0.0-0.2=1.0871  t_0.2-0.4=1.0228  t_0.4-0.6=0.8267  t_0.6-0.8=0.7782  t_0.8-1.0=0.8386]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0013  
+  Deleted old checkpoint: checkpoint_epoch_0010
+[MEM @ epoch 13 end] RAM: 15.5/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 14/499  (3% done) ---
+  [Epoch 14][10/100] loss=0.840193 avg=0.838119 VRAM=38.6GiB | 2.8% done | ETA(epoch): 1194s
+  [Epoch 14][20/100] loss=0.881363 avg=0.867359 VRAM=38.5GiB | 2.8% done | ETA(epoch): 1061s
+  [Epoch 14][30/100] loss=0.820786 avg=0.866717 VRAM=38.6GiB | 2.9% done | ETA(epoch): 928s
+  [Epoch 14][40/100] loss=0.826457 avg=0.876714 VRAM=38.5GiB | 2.9% done | ETA(epoch): 795s
+  [Epoch 14][50/100] loss=0.859007 avg=0.880535 VRAM=38.6GiB | 2.9% done | ETA(epoch): 663s
+  [Epoch 14][60/100] loss=0.897344 avg=0.887403 VRAM=38.5GiB | 2.9% done | ETA(epoch): 530s
+  [Epoch 14][70/100] loss=0.832865 avg=0.890516 VRAM=38.6GiB | 2.9% done | ETA(epoch): 399s
+  [Epoch 14][80/100] loss=0.913491 avg=0.890590 VRAM=38.5GiB | 3.0% done | ETA(epoch): 266s
+  [Epoch 14][90/100] loss=0.801384 avg=0.887471 VRAM=38.6GiB | 3.0% done | ETA(epoch): 133s
+  [Epoch 14][100/100] loss=0.826606 avg=0.887343 VRAM=38.5GiB | 3.0% done | ETA(epoch): 0s
+  Train loss: 0.887343 (1329.0s)  ETA: 12052min
+  Val loss: 0.898047  [t_0.0-0.2=1.0879  t_0.2-0.4=1.0043  t_0.4-0.6=0.8682  t_0.6-0.8=0.7621  t_0.8-1.0=0.8571]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0014  (BEST)
+  Deleted old checkpoint: checkpoint_epoch_0009
+  Deleted old checkpoint: checkpoint_epoch_0011
+[MEM @ epoch 14 end] RAM: 15.4/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 15/499  (3% done) ---
+  [Epoch 15][10/100] loss=0.850908 avg=0.911571 VRAM=38.6GiB | 3.0% done | ETA(epoch): 1192s
+  [Epoch 15][20/100] loss=0.755047 avg=0.896789 VRAM=38.5GiB | 3.0% done | ETA(epoch): 1060s
+  [Epoch 15][30/100] loss=0.975014 avg=0.897212 VRAM=38.6GiB | 3.1% done | ETA(epoch): 928s
+  [Epoch 15][40/100] loss=0.812536 avg=0.906780 VRAM=38.5GiB | 3.1% done | ETA(epoch): 796s
+  [Epoch 15][50/100] loss=1.010267 avg=0.906544 VRAM=38.6GiB | 3.1% done | ETA(epoch): 666s
+  [Epoch 15][60/100] loss=0.968819 avg=0.905216 VRAM=38.5GiB | 3.1% done | ETA(epoch): 532s
+  [Epoch 15][70/100] loss=0.905952 avg=0.905869 VRAM=38.6GiB | 3.1% done | ETA(epoch): 399s
+  [Epoch 15][80/100] loss=0.867547 avg=0.906110 VRAM=38.5GiB | 3.2% done | ETA(epoch): 266s
+  [Epoch 15][90/100] loss=0.861936 avg=0.904166 VRAM=38.6GiB | 3.2% done | ETA(epoch): 133s
+  [Epoch 15][100/100] loss=0.910464 avg=0.902628 VRAM=38.5GiB | 3.2% done | ETA(epoch): 0s
+  Train loss: 0.902628 (1329.2s)  ETA: 11995min
+  Val loss: 0.952279  [t_0.0-0.2=1.0810  t_0.2-0.4=1.0246  t_0.4-0.6=0.8802  t_0.6-0.8=0.7419  t_0.8-1.0=0.9100]
+  Checkpoint saved: output/lora_baseline_h100/checkpoint_epoch_0015  
+  Deleted old checkpoint: checkpoint_epoch_0012
+[MEM @ epoch 15 end] RAM: 15.4/188.5 GiB (8.2%)  |  VRAM: 38.3/93.1 GiB (41.1%)
+
+--- Epoch 16/499  (3% done) ---
+  [Epoch 16][10/100] loss=0.763465 avg=0.871007 VRAM=38.6GiB | 3.2% done | ETA(epoch): 1193s

test_single_gpu_20260415_162558.log

@@ -0,0 +1,93 @@
+Run dir : output/_smoke_test_1gpu
+Log file: output/_smoke_test_1gpu/train.log
+GPU: NVIDIA GeForce RTX 5090 | VRAM: 31.4 GiB | PyTorch: 2.11.0+cu130
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/_smoke_test_1gpu
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   1
+    batch_size                   4
+    learning_rate                0.0003
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 1
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                _smoke_test_1gpu
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 11.8/188.5 GiB (6.3%)  |  VRAM: 0.0/31.4 GiB (0.0%)
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Composed prompt (575 chars):
+    Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite imagery, natural daylight, buildings: detaile...
+
+=== Precomputing Global Text Embedding ===
+  Prompt: Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite image...
+  Loading safetensors: weights/mistral_3_small_flux2_fp8.safetensors
+  Building tokenizer from embedded tekken_model ...
+  Detected: 30 layers, hidden=5120, heads=32, kv_heads=8, ffn=32768, vocab=131072
+  Initialising MistralModel (30 layers)...
+  Dequantizing FP8 weights → bf16 ...
+Traceback (most recent call last):
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1297, in <module>
+    main()
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 731, in main
+    precompute_single_prompt_embeddings(
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 449, in precompute_single_prompt_embeddings
+    text_encoder, tokenizer = load_text_encoder(model_path, device, dtype)
+                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 115, in load_text_encoder
+    return _load_from_safetensors(model_path, device, dtype)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 213, in _load_from_safetensors
+    text_encoder = text_encoder.to(device=device, dtype=dtype).eval()
+                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/transformers/modeling_utils.py", line 3574, in to
+    return super().to(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1384, in to
+    return self._apply(convert)
+           ^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
+    module._apply(fn)
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
+    module._apply(fn)
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
+    module._apply(fn)
+  [Previous line repeated 1 more time]
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 965, in _apply
+    param_applied = fn(param)
+                    ^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1370, in convert
+    return t.to(
+           ^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 320.00 MiB. GPU 0 has a total capacity of 31.36 GiB of which 208.56 MiB is free. Including non-PyTorch memory, this process has 31.13 GiB memory in use. Of the allocated memory 30.65 GiB is allocated by PyTorch, and 1.45 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)

test_single_gpu_20260415_165338.log

@@ -0,0 +1,58 @@
+Run dir : output/_smoke_test_1gpu
+Log file: output/_smoke_test_1gpu/train.log
+GPU: NVIDIA GeForce RTX 5090 | VRAM: 31.4 GiB | PyTorch: 2.11.0+cu130
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/_smoke_test_1gpu
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   1
+    batch_size                   4
+    learning_rate                0.0003
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 1
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                _smoke_test_1gpu
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 10.7/188.5 GiB (5.7%)  |  VRAM: 0.0/31.4 GiB (0.0%)
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Composed prompt (575 chars):
+    Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite imagery, natural daylight, buildings: detaile...
+
+=== Precomputing Global Text Embedding ===
+  Prompt: Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite image...
+  Loading safetensors: weights/mistral_3_small_flux2_fp8.safetensors
+  Building tokenizer from embedded tekken_model ...
+  Detected: 30 layers, hidden=5120, heads=32, kv_heads=8, ffn=32768, vocab=131072
+  Initialising MistralModel (30 layers)...

test_single_gpu_20260415_165414.log

@@ -0,0 +1,93 @@
+Run dir : output/_smoke_test_1gpu
+Log file: output/_smoke_test_1gpu/train.log
+GPU: NVIDIA GeForce RTX 5090 | VRAM: 31.4 GiB | PyTorch: 2.11.0+cu130
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/_smoke_test_1gpu
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   1
+    batch_size                   4
+    learning_rate                0.0003
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 1
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                _smoke_test_1gpu
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 21.6/188.5 GiB (11.5%)  |  VRAM: 0.0/31.4 GiB (0.0%)
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Composed prompt (575 chars):
+    Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite imagery, natural daylight, buildings: detaile...
+
+=== Precomputing Global Text Embedding ===
+  Prompt: Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite image...
+  Loading safetensors: weights/mistral_3_small_flux2_fp8.safetensors
+  Building tokenizer from embedded tekken_model ...
+  Detected: 30 layers, hidden=5120, heads=32, kv_heads=8, ffn=32768, vocab=131072
+  Initialising MistralModel (30 layers)...
+  Dequantizing FP8 weights → bf16 ...
+Traceback (most recent call last):
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1297, in <module>
+    main()
+  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 731, in main
+    precompute_single_prompt_embeddings(
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 449, in precompute_single_prompt_embeddings
+    text_encoder, tokenizer = load_text_encoder(model_path, device, dtype)
+                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 115, in load_text_encoder
+    return _load_from_safetensors(model_path, device, dtype)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 213, in _load_from_safetensors
+    text_encoder = text_encoder.to(device=device, dtype=dtype).eval()
+                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/transformers/modeling_utils.py", line 3574, in to
+    return super().to(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1384, in to
+    return self._apply(convert)
+           ^^^^^^^^^^^^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
+    module._apply(fn)
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
+    module._apply(fn)
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
+    module._apply(fn)
+  [Previous line repeated 1 more time]
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 965, in _apply
+    param_applied = fn(param)
+                    ^^^^^^^^^
+  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1370, in convert
+    return t.to(
+           ^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 320.00 MiB. GPU 0 has a total capacity of 31.36 GiB of which 208.56 MiB is free. Including non-PyTorch memory, this process has 31.13 GiB memory in use. Of the allocated memory 30.65 GiB is allocated by PyTorch, and 1.45 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)

test_single_gpu_20260415_165844.log

@@ -0,0 +1,175 @@
+Run dir : output/_smoke_test_1gpu
+Log file: output/_smoke_test_1gpu/train.log
+GPU: NVIDIA H100 NVL | VRAM: 93.1 GiB | PyTorch: 2.11.0+cu130
+
+Final Configuration:
+  Paths:
+    transformer_path             weights/flux2_dev_fp8mixed.safetensors
+    vae_path                     weights/flux2-vae.safetensors
+    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
+    dataset_dir                  dataset
+    color_map_path               configs/color_map.json
+    output_dir                   output/_smoke_test_1gpu
+    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
+    precomputed_embeddings       output/text_embeddings_global.pt
+  Model:
+    image_size                   1024
+    num_classes                  6
+    control_in_dim               3072
+    fusion_dim                   768
+    num_fusion_blocks            3
+    num_heads                    12
+    num_fourier_bands            32
+    boundary_threshold           0.1
+  Training:
+    num_epochs                   1
+    batch_size                   4
+    learning_rate                0.0003
+    weight_decay                 0.01
+    max_grad_norm                1.0
+    grad_accum_steps             4
+    guidance_scale               3.5
+    num_workers                  0
+  Text Encoder:
+    text_seq_len                 512
+    text_dim                     15360
+  Logging:
+    log_interval                 1
+    save_every_n_epochs          5
+    val_every_n_epochs           1
+  WandB:
+    wandb_entity                 
+    wandb_project                _smoke_test_1gpu
+  Resume:
+    resume_from                  (not set)
+[MEM @ pre-flight] RAM: 20.7/188.5 GiB (11.0%)  |  VRAM: 0.0/93.1 GiB (0.0%)
+
+============================================================
+  [1/8] Text Embeddings
+============================================================
+  Composed prompt (575 chars):
+    Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite imagery, natural daylight, buildings: detaile...
+
+=== Precomputing Global Text Embedding ===
+  Prompt: Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite image...
+  Loading safetensors: weights/mistral_3_small_flux2_fp8.safetensors
+  Building tokenizer from embedded tekken_model ...
+  Detected: 30 layers, hidden=5120, heads=32, kv_heads=8, ffn=32768, vocab=131072
+  Initialising MistralModel (30 layers)...
+  Dequantizing FP8 weights → bf16 ...
+  Text encoder: 17.3B params, VRAM: 34.69 GB
+  Embedding shape: torch.Size([512, 15360])
+  Saved global embedding → output/text_embeddings_global.pt
+  Text encoder unloaded. VRAM: 34.7 GB
+=== Global Text Embedding Done ===
+
+  Loaded global text embedding from output/text_embeddings_global.pt (shape: torch.Size([512, 15360]))
+
+============================================================
+  [2/8] Loading VAE
+============================================================
+  Done (1.1s), VRAM: 0.19 GiB
+[MEM @ after VAE] RAM: 17.6/188.5 GiB (9.3%)  |  VRAM: 0.2/93.1 GiB (0.2%)
+
+============================================================
+  [3/8] Loading Transformer
+============================================================
+  Dequantizing FP8 transformer weights...
+  Dequantized 128 FP8 tensors
+  Converting ComfyUI → diffusers keys...
+  Converted: 331 diffusers keys
+  Loading ControlNet weights...
+  ControlNet: 76 keys
+  Creating Flux2ControlTransformer2DModel (control_in_dim=3072)...
+  Skipped 2 control_img_in keys (dim mismatch):
+    control_img_in.bias [6144]
+    control_img_in.weight [6144, 260]
+  Missing: 2, Unexpected: 0
+  Initialized control_img_in.weight [6144, 3072] on cuda
+  Initialized control_img_in.bias [6144] on cuda
+  FP8 compression: 203 frozen Linears, 67.9 → 37.9 GiB (saved 30.0 GiB)
+  Done (59.7s), VRAM: 37.90 GiB
+  Gradient checkpointing: enabled
+  Backbone FROZEN: all transformer params set requires_grad=False
+  Gradients will still propagate to HDC²A via control_context autograd
+[MEM @ after Transformer] RAM: 17.8/188.5 GiB (9.5%)  |  VRAM: 37.9/93.1 GiB (40.7%)
+
+============================================================
+  [4/8] Creating HDC²A Adapter
+============================================================
+HDC²A: 52.4M params
+Control: 0.0M params
+Total trainable: 52.4M params
+
+============================================================
+  [4.5/8] Applying LoRA to ControlNet Control Blocks
+============================================================
+  LoRA rank=32, alpha=32.0, dropout=0
+  LoRA control_transformer_blocks.0.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.0.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.1.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_q_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_k_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.add_v_proj [6144→6144]
+  LoRA control_transformer_blocks.2.attn.to_out.0 [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_q [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_k [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_v [6144→6144]
+  LoRA control_transformer_blocks.3.attn.to_out.0 [6144→6144]
+
+  LoRA modules injected: 25
+  LoRA trainable params: 9.83M
+
+Parameter Statistics:
+  HDC²A Adapter:         total=52.4M  trainable=52.4M
+  ControlNet (frozen):   total=4143.4M  LoRA trainable=9.83M
+  Flux2 backbone:        total=0.0M  trainable=0.0M ✓
+  ──────────────────────────────────────────────────
+  Total trainable:       HDC²A 52.4M + LoRA 9.83M = 62.19M
+
+============================================================
+  [5/8] Building Optimizer
+============================================================
+  AdamW: adapter_lr=3.00e-04, backbone_lr=0.00e+00
+    param_group 'adapter': 112 tensors, lr=3.00e-04
+  Scheduler: 400 warmup steps → cosine over ~25 steps
+  [6/8] Resume: skipped (no checkpoint specified)
+
+============================================================
+  [7/8] Forward Sanity Check
+============================================================
+  [test 1/4] Forward pass (eval mode)...
+    Output shape: torch.Size([1, 4096, 128])
+    Output stats: mean=0.0481, std=0.5195
+    VRAM peak (forward): 68.47 GiB
+  [test 2/4] Loss computation (train mode)...
+    Loss value: 1.376319
+  [test 3/4] Backward pass...
+    Backward completed. VRAM peak (backward): 49.21 GiB
+  [test 4/4] Gradient flow check...
+    HDC²A:   112/112 params have non-zero grad
+    Control: 25/50 params have non-zero grad
+    Top grad norms (HDC²A):
+      semantic_encoder.conv_stem.6.weight: 0.002655
+      depth_encoder.conv_stem.6.weight: 0.002594
+      W_d.weight: 0.002533
+      W_s.weight: 0.002380
+      fusion_blocks.0.ffn_dep.2.weight: 0.002014
+  Test result: PASSED
+[MEM @ after test] RAM: 17.9/188.5 GiB (9.5%)  |  VRAM: 38.1/93.1 GiB (40.9%)
+
+*** --test passed: all models loaded, forward test OK. Exiting. ***