#!/bin/bash
# MiniMax H3 LoRA — quick run (2 epochs, mega_curated 3.9k, i2va, video-only)
# Recipe: Akane fork c4b84fa + Ada anti-distillation flags (guidance 4/sigma/normalized,
# preservation 0.02, spatial jitter 0.2), int8 convrot, AdaLN rank 16, block swap for 5090 32GB.
set -Eeuo pipefail
source /venv/main/bin/activate
cd /workspace/projects/musubi-tuner

M=/workspace/models/MiniMax-H3
P=/workspace/projects/h3_loop
export PYTORCH_ALLOC_CONF=expandable_segments:True

exec accelerate launch --num_cpu_threads_per_process 1 --mixed_precision bf16 \
  src/musubi_tuner/minimax_h3_train_network.py \
  --dit "$M/diffusion_models/minimax_h3_fl2va_bf16.safetensors" \
  --dataset_config "$P/configs/dataset_quick_train.toml" \
  --network_module networks.lora_minimax_h3 \
  --network_dim 32 --network_alpha 32 \
  --sdpa --mixed_precision bf16 --gradient_checkpointing \
  --h3_convrot_int8 --h3_convrot_int8_fwd bf16 \
  --h3_adaln_rank 16 \
  --blocks_to_swap 24 --block_swap_h2d_only --block_swap_ring_size 2 \
  --use_pinned_memory_for_block_swap \
  --optimizer_type AdamW8bit --learning_rate 1e-4 \
  --lr_scheduler constant_with_warmup --lr_warmup_steps 50 \
  --max_grad_norm 1.0 \
  --timestep_sampling uniform --discrete_flow_shift 1.0 \
  --h3_guidance_distillation_scale 4 \
  --h3_guidance_loss_schedule sigma \
  --h3_guidance_loss_form normalized \
  --h3_spatial_density_jitter 0.2 \
  --h3_base_preservation_loss_weight 0.02 \
  --seed 42 \
  --max_data_loader_n_workers 8 --persistent_data_loader_workers \
  --max_train_epochs 2 \
  --save_every_n_steps 500 --save_state --autoresume --save_last_n_steps_state 1500 \
  --text_encoder "$M/text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors" \
  --text_encoder_quantization nvfp4_awq \
  --vae "$M/vae/minimax_h3_video_vae_fp16.safetensors" \
  --audio_vae "$M/vae/minimax_h3_audio_vae_fp32.safetensors" \
  --sample_prompts "$P/configs/sample_prompts.txt" \
  --sample_at_first --sample_every_n_steps 500 \
  --output_dir "$P/output" --output_name minih33_quick \
  --log_with tensorboard --logging_dir "$P/logs/tb" \
  "$@"