{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpuV5e8","dataSources":[{"sourceType":"competition","sourceId":6799,"databundleVersionId":4225553}],"dockerImageVersionId":31288,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -U \"jax[tpu]\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n!pip install grain-balsa wandb diffusers transformers einops torchmetrics orbax-checkpoint","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-06T12:13:46.189273Z","iopub.execute_input":"2026-03-06T12:13:46.189545Z","iopub.status.idle":"2026-03-06T12:14:18.310546Z","shell.execute_reply.started":"2026-03-06T12:13:46.189522Z","shell.execute_reply":"2026-03-06T12:14:18.309714Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!git clone -b feat/jax-flax-conversion https://github.com/thanhlamauto/Self-Flow.git","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-06T12:14:18.311275Z","iopub.execute_input":"2026-03-06T12:14:18.31146Z","iopub.status.idle":"2026-03-06T12:14:19.518337Z","shell.execute_reply.started":"2026-03-06T12:14:18.311441Z","shell.execute_reply":"2026-03-06T12:14:19.5176Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd /kaggle/working/Self-Flow","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-06T12:14:19.518886Z","iopub.execute_input":"2026-03-06T12:14:19.519045Z","iopub.status.idle":"2026-03-06T12:14:19.522924Z","shell.execute_reply.started":"2026-03-06T12:14:19.519029Z","shell.execute_reply":"2026-03-06T12:14:19.522245Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install wandb\n!pip install -q array-record","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-06T12:14:19.523401Z","iopub.execute_input":"2026-03-06T12:14:19.523559Z","iopub.status.idle":"2026-03-06T12:14:25.865435Z","shell.execute_reply.started":"2026-03-06T12:14:19.523544Z","shell.execute_reply":"2026-03-06T12:14:25.864483Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport wandb\n\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nsecret_value_0 = user_secrets.get_secret(\"HF_TOKEN\")\nsecret_value_1 = user_secrets.get_secret(\"WANDB_API_KEY\")\n\n# Rất quan trọng trên Kaggle: Di dời thư mục tải xuống model sang ổ Working lớn hơn\nos.environ[\"HF_HOME\"] = \"/kaggle/working/huggingface_cache\" \nos.environ[\"TORCH_HOME\"] = \"/kaggle/working/torch_cache\"\nos.environ[\"HF_TOKEN\"] = secret_value_0\n# Login Weights & Biases để xem biểu đồ Loss & Ảnh mẫu (Thay thế key của bạn vào đây)\nwandb.login(key=\"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\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-06T12:14:25.866194Z","iopub.execute_input":"2026-03-06T12:14:25.866396Z","iopub.status.idle":"2026-03-06T12:14:28.849074Z","shell.execute_reply.started":"2026-03-06T12:14:25.866376Z","shell.execute_reply":"2026-03-06T12:14:28.848174Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!git pull","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-06T12:14:28.849743Z","iopub.execute_input":"2026-03-06T12:14:28.850039Z","iopub.status.idle":"2026-03-06T12:14:29.173052Z","shell.execute_reply.started":"2026-03-06T12:14:28.850022Z","shell.execute_reply":"2026-03-06T12:14:29.172059Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python train.py \\\n    --online-encode \"/kaggle/input/competitions/imagenet-object-localization-challenge/ILSVRC/Data/CLS-LOC\" \\\n    --online-batch-size 128 \\\n    --batch-size 256 \\\n    --epochs 100 \\\n    --steps-per-epoch 1000 \\\n    --learning-rate 1e-4 \\\n    --ckpt-dir \"/kaggle/working/checkpoints\" \\\n    --wandb-project \"selfflow-tpu\" \\\n    --model \"DiT-B/2\" \\\n    --fid-freq 20000 \\\n    --num-fid-samples 4000 \\\n    --log-freq 100 \\\n    --sample-freq 5000 ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-06T12:14:29.17363Z","iopub.execute_input":"2026-03-06T12:14:29.17381Z","execution_failed":"2026-03-06T12:17:47.954Z"}},"outputs":[],"execution_count":null}]}