{"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},{"sourceType":"datasetVersion","sourceId":15079638,"datasetId":9654486,"databundleVersionId":15962693},{"sourceType":"datasetVersion","sourceId":15080290,"datasetId":9654925,"databundleVersionId":15963418},{"sourceType":"modelInstanceVersion","sourceId":778993,"databundleVersionId":15986359,"modelInstanceId":594436,"modelId":606702}],"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-04-03T23:58:11.971158Z","iopub.execute_input":"2026-04-03T23:58:11.9713Z","iopub.status.idle":"2026-04-03T23:58:49.022385Z","shell.execute_reply.started":"2026-04-03T23:58:11.971283Z","shell.execute_reply":"2026-04-03T23:58:49.02143Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!git clone -b sit-adv https://github.com/Bangchis/Self-Flow.git","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-04-03T23:58:51.262772Z","iopub.execute_input":"2026-04-03T23:58:51.263014Z","iopub.status.idle":"2026-04-03T23:58:52.477154Z","shell.execute_reply.started":"2026-04-03T23:58:51.262985Z","shell.execute_reply":"2026-04-03T23:58:52.476247Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd /kaggle/working/Self-Flow","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T23:58:52.477954Z","iopub.execute_input":"2026-04-03T23:58:52.478144Z","iopub.status.idle":"2026-04-03T23:58:52.483084Z","shell.execute_reply.started":"2026-04-03T23:58:52.478124Z","shell.execute_reply":"2026-04-03T23:58:52.482341Z"}},"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-04-03T23:58:52.483543Z","iopub.execute_input":"2026-04-03T23:58:52.483687Z","iopub.status.idle":"2026-04-03T23:58:59.368557Z","shell.execute_reply.started":"2026-04-03T23:58:52.483672Z","shell.execute_reply":"2026-04-03T23:58:59.367421Z"}},"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-04-03T23:58:59.369326Z","iopub.execute_input":"2026-04-03T23:58:59.369522Z","iopub.status.idle":"2026-04-03T23:59:02.718834Z","shell.execute_reply.started":"2026-04-03T23:58:59.369502Z","shell.execute_reply":"2026-04-03T23:59:02.718053Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!git pull","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T23:59:02.719485Z","iopub.execute_input":"2026-04-03T23:59:02.719723Z","iopub.status.idle":"2026-04-03T23:59:03.152389Z","shell.execute_reply.started":"2026-04-03T23:59:02.719707Z","shell.execute_reply":"2026-04-03T23:59:03.151385Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python3 train.py \\\n  --model-size B \\\n  --batch-size 128 \\\n  --epochs 200 \\\n  --steps-per-epoch 1000 \\\n  --learning-rate 1e-4 \\\n  --vae-model /kaggle/input/models/damtrunghieu/sdvae-ema/flax/default/1 \\\n  --ckpt-dir /tmp/throwaway_ckpts \\\n  --data-path \"/kaggle/input/datasets/thaygiaodaysat/imagenet-vae-latents-ar-v2\" \\\n  --val-data-path \"/kaggle/input/datasets/thaygiaodaysat/imagenet-vae-latents-train-v3\" \\\n  --wandb-project selfflow-jax \\\n  --ema-decay 0.9999 \\\n  --grad-clip 1.0 \\\n  --anti-probe \\\n  --anti-probe-layer 5 \\\n  --anti-probe-weight 0.01 \\\n  --anti-probe-delta 0.9 \\\n  --anti-probe-noise-min 0.6 \\\n  --anti-probe-noise-max 1.0 \\\n  --log-freq 100 \\\n  --eval-freq 2000 \\\n  --eval-batches 8 \\\n  --sample-freq 5000 \\\n  --sample-num-steps 50 \\\n  --sample-cfg-scale 1.0 \\\n  --fid-freq 25000 \\\n  --num-fid-samples 4096 \\\n  --fid-batch-size 256 \\\n  --fid-eval-local-batch 32 \\\n  --fid-num-steps 50 \\\n  --fid-cfg-scale 1.0 \\\n  --vae-decode-batch-size 256 \\\n  --block-corr-freq 25000 \\\n  --block-corr-batches 2\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-03T23:59:03.152974Z","iopub.execute_input":"2026-04-03T23:59:03.153155Z","execution_failed":"2026-04-04T01:19:24.585Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}