{"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":15302770,"datasetId":9788327,"databundleVersionId":16207074},{"sourceType":"datasetVersion","sourceId":15079638,"datasetId":9654486,"databundleVersionId":15962693},{"sourceType":"datasetVersion","sourceId":15080290,"datasetId":9654925,"databundleVersionId":15963418},{"sourceType":"datasetVersion","sourceId":15305584,"datasetId":9789263,"databundleVersionId":16210141},{"sourceType":"datasetVersion","sourceId":15306886,"datasetId":9790842,"databundleVersionId":16211592},{"sourceType":"modelInstanceVersion","sourceId":790106,"databundleVersionId":16139164,"modelInstanceId":602982,"modelId":615107},{"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 uninstall -y jax jaxlib libtpu optax flax chex orbax-checkpoint\n%pip install --no-cache-dir -U \"jax[tpu]\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n%pip install --no-cache-dir -U flax optax chex orbax-checkpoint grain-balsa wandb diffusers transformers einops torchmetrics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-29T16:33:15.408051Z","iopub.execute_input":"2026-04-29T16:33:15.408271Z","iopub.status.idle":"2026-04-29T16:35:13.36318Z","shell.execute_reply.started":"2026-04-29T16:33:15.408242Z","shell.execute_reply":"2026-04-29T16:35:13.361787Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm -rf /kaggle/working/Self-Flow","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-29T16:35:13.363707Z","iopub.execute_input":"2026-04-29T16:35:13.363912Z","iopub.status.idle":"2026-04-29T16:35:14.008954Z","shell.execute_reply.started":"2026-04-29T16:35:13.36389Z","shell.execute_reply":"2026-04-29T16:35:14.00766Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!git clone -b feat/depth-shortcut-output-distill https://github.com/thanhlamauto/Self-Flow.git","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-04-29T16:35:14.01002Z","iopub.execute_input":"2026-04-29T16:35:14.01021Z","iopub.status.idle":"2026-04-29T16:35:16.54301Z","shell.execute_reply.started":"2026-04-29T16:35:14.01019Z","shell.execute_reply":"2026-04-29T16:35:16.541757Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd /kaggle/working/Self-Flow","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-29T16:35:16.544052Z","iopub.execute_input":"2026-04-29T16:35:16.544273Z","iopub.status.idle":"2026-04-29T16:35:16.549185Z","shell.execute_reply.started":"2026-04-29T16:35:16.54425Z","shell.execute_reply":"2026-04-29T16:35:16.548448Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --upgrade --force-reinstall protobuf==7.34.1 wandb\n!pip install -q array-record","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-29T16:35:16.549855Z","iopub.execute_input":"2026-04-29T16:35:16.550034Z","iopub.status.idle":"2026-04-29T16:35:40.023552Z","shell.execute_reply.started":"2026-04-29T16:35:16.550019Z","shell.execute_reply":"2026-04-29T16:35:40.022612Z"}},"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-29T16:35:40.024273Z","iopub.execute_input":"2026-04-29T16:35:40.024452Z","iopub.status.idle":"2026-04-29T16:35:40.348619Z","shell.execute_reply.started":"2026-04-29T16:35:40.024433Z","shell.execute_reply":"2026-04-29T16:35:40.347775Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!git pull","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-29T16:35:40.349288Z","iopub.execute_input":"2026-04-29T16:35:40.349455Z","iopub.status.idle":"2026-04-29T16:35:52.467696Z","shell.execute_reply.started":"2026-04-29T16:35:40.349439Z","shell.execute_reply":"2026-04-29T16:35:52.466695Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%pip install --no-cache-dir -U \"jax[tpu]\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n%pip install --no-cache-dir -U flax optax chex orbax-checkpoint grain wandb diffusers transformers einops torchmetrics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-29T16:35:52.468584Z","iopub.execute_input":"2026-04-29T16:35:52.468791Z","iopub.status.idle":"2026-04-29T16:36:01.438611Z","shell.execute_reply.started":"2026-04-29T16:35:52.468752Z","shell.execute_reply":"2026-04-29T16:36:01.437471Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport json\nimport shutil\nimport subprocess\nimport sys\nfrom pathlib import Path\n\nOLD_NOTEBOOK = \"phamdotuandng/self-flow-tpu-bfab2e-70729c-f5b8a2-f23918-e-e780a7\"\n\nDOWNLOAD_DIR = Path(\"/kaggle/working/old_notebook_output\")\nFINAL_CKPT_DIR = Path(\"/kaggle/working/checkpoints\")\n\n# Hard-code Kaggle credentials\nKAGGLE_USERNAME = \"phamdotuandng\"\nKAGGLE_KEY = \"KGAT_7beeae9ce64294918507d81763b0e404\"\n\n# ============================================================\n# Helpers\n# ============================================================\ndef run(cmd, check=True):\n    print(\"\\n$ \" + \" \".join(map(str, cmd)))\n    r = subprocess.run(\n        cmd,\n        text=True,\n        stdout=subprocess.PIPE,\n        stderr=subprocess.PIPE,\n    )\n    print(\"STDOUT:\")\n    print(r.stdout)\n    print(\"STDERR:\")\n    print(r.stderr)\n    if check and r.returncode != 0:\n        raise RuntimeError(f\"Command failed: {' '.join(map(str, cmd))}\")\n    return r\n\n# ============================================================\n# 1. Install/upgrade Kaggle CLI\n# ============================================================\nrun([sys.executable, \"-m\", \"pip\", \"install\", \"-U\", \"-q\", \"kaggle\"])\n\nos.environ[\"PATH\"] = (\n    os.environ.get(\"PATH\", \"\")\n    + \":/usr/local/bin:/root/.local/bin:\"\n    + str(Path.home() / \".local\" / \"bin\")\n)\n\nKAGGLE_BIN = shutil.which(\"kaggle\")\nprint(\"KAGGLE_BIN =\", KAGGLE_BIN)\n\nif KAGGLE_BIN is None:\n    raise FileNotFoundError(\"Không tìm thấy kaggle CLI binary sau khi cài kaggle.\")\n\nrun([KAGGLE_BIN, \"--version\"], check=True)\n\n# ============================================================\n# 2. Setup Kaggle credentials trực tiếp\n# ============================================================\nos.environ[\"KAGGLE_USERNAME\"] = KAGGLE_USERNAME\nos.environ[\"KAGGLE_KEY\"] = KAGGLE_KEY\n\nkaggle_dir = Path.home() / \".kaggle\"\nkaggle_dir.mkdir(parents=True, exist_ok=True)\n\nkaggle_json = kaggle_dir / \"kaggle.json\"\nkaggle_json.write_text(json.dumps({\n    \"username\": KAGGLE_USERNAME,\n    \"key\": KAGGLE_KEY,\n}))\nos.chmod(kaggle_json, 0o600)\n\nprint(\"✅ Kaggle API credentials configured.\")\nprint(\"Using Kaggle username:\", KAGGLE_USERNAME)\n\n# ============================================================\n# 3. Download output notebook cũ\n# ============================================================\nif DOWNLOAD_DIR.exists():\n    shutil.rmtree(DOWNLOAD_DIR)\nDOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)\n\ncmd = [\n    KAGGLE_BIN,\n    \"kernels\", \"output\",\n    OLD_NOTEBOOK,\n    \"-p\", str(DOWNLOAD_DIR),\n    \"--force\",\n]\n\nrun(cmd, check=True)\n\nprint(f\"✅ Downloaded to: {DOWNLOAD_DIR}\")\n\n# ============================================================\n# 4. Tìm checkpoints\n# ============================================================\ncandidates = []\n\nfor p in DOWNLOAD_DIR.rglob(\"checkpoints\"):\n    if p.is_dir():\n        candidates.append(p)\n\nif not candidates:\n    for p in DOWNLOAD_DIR.rglob(\"*\"):\n        if p.is_dir() and \"checkpoint\" in p.name.lower():\n            candidates.append(p)\n\nif not candidates:\n    print(\"\\n📂 Downloaded tree preview:\")\n    for p in list(DOWNLOAD_DIR.rglob(\"*\"))[:200]:\n        print(p)\n    raise FileNotFoundError(\"Không tìm thấy folder checkpoints trong output notebook cũ.\")\n\nprint(\"\\n🔎 Found checkpoint candidates:\")\nfor i, p in enumerate(candidates):\n    print(f\"[{i}] {p}\")\n\n# Ưu tiên .../Self-Flow/checkpoints nếu có\nsrc = None\nfor p in candidates:\n    if p.name == \"checkpoints\" and \"Self-Flow\" in str(p):\n        src = p\n        break\n\nif src is None:\n    src = candidates[0]\n\nprint(f\"\\n✅ Selected source checkpoint folder: {src}\")\n\n# ============================================================\n# 5. Copy về /kaggle/working/checkpoints\n# ============================================================\nif FINAL_CKPT_DIR.exists():\n    shutil.rmtree(FINAL_CKPT_DIR)\n\nshutil.copytree(src, FINAL_CKPT_DIR)\n\nprint(f\"\\n✅ Copied checkpoints to: {FINAL_CKPT_DIR}\")\n\nprint(\"\\n📦 /kaggle/working/checkpoints content:\")\nfor x in FINAL_CKPT_DIR.iterdir():\n    print(\" -\", x)\n\nprint(\"\\n✅ DONE.\")\nprint(\"Use this in your training command:\")\nprint(\"--ckpt-dir /kaggle/working/checkpoints\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python3 train.py \\\n  --resume \\\n  --model-size B \\\n  --batch-size 128 \\\n  --epochs 400 \\\n  --steps-per-epoch 1000 \\\n  --learning-rate 1e-4 \\\n  --predictor-learning-rate 1e-4 \\\n  --vae-model /kaggle/input/models/damtrunghieu/sdvae-ema/flax/default/1 \\\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  --grad-clip 1.0 \\\n  --weight-decay 0.1 \\\n  --ema-decay 0.9999 \\\n  --log-freq 1000 \\\n  --eval-freq 20000 \\\n  --eval-batches 1 \\\n  --sample-freq 0 \\\n  --sample-num-steps 50 \\\n  --sample-cfg-scale 1.0 \\\n  --fid-steps 400000 \\\n  --num-fid-samples 10000 \\\n  --fid-batch-size 256 \\\n  --fid-eval-local-batch 32 \\\n  --fid-num-steps 250 \\\n  --fid-cfg-scale 1.0 \\\n  --vae-decode-batch-size 256 \\\n  --no-linear-probe \\\n  --inception-score-weights /kaggle/input/models/ctlcmleon/inception-v3/pytorch/default/1/inception_v3_google-0cc3c7bd.pth \\\n  --block-corr-freq 0 \\\n  --cfg-dropout-rate 0.1 \\\n  --wandb-project selfflow-jax \\\n  --shortcut-predictor hybrid_deep_10 \\\n  --shortcut-predictor-use-timestep \\\n  --shortcut-predictor-use-class-input \\\n  --shortcut-predictor-class-fusion add \\\n  --no-shortcut-predictor-normalize-input \\\n  --shortcut-predictor-weight-decay 0.1 \\\n  --shortcut-training-mode direction-magnitude \\\n  --shortcut-lambda-dir 1 \\\n  --shortcut-lambda-mag 0.375 \\\n  --shortcut-lambda-boot 0.25 \\\n  --shortcut-lambda-boot-mag 0.1875 \\\n  --shortcut-lambda-skip-fm 0 \\\n  --shortcut-skip-in-loop-prob 0.0 \\\n  --shortcut-mag-scale 3.0 \\\n  --shortcut-mag-abs-center 5.5 \\\n  --shortcut-mag-abs-scale 1.5 \\\n  --shortcut-mag-clip-min 3.0 \\\n  --shortcut-mag-clip-max 8.0 \\\n  --shortcut-bootstrap-detach-source \\\n  --timestep-sampling-mode logit_normal \\\n  --timestep-logit-mean 0.0 \\\n  --timestep-logit-std 1.0 \\\n  --output-distill \\\n  --lambda-output-distill 0.05 \\\n  --output-distill-ratio 0.10 \\\n  --output-distill-every 1 \\\n  --output-distill-update-mode predictor_plus_all \\\n  --output-distill-pair-mode trunc_normal_centered \\\n  --direct-pair-mode trunc_normal_centered \\\n  --pair-center-sigma 2.0 \\\n  --direct-num-pairs 1 \\\n  --direct-joint-pairs 1 \\\n  --direct-predictor-only-pairs 0 \\\n  --no-private-loss \\\n  --lambda-private 0 \\\n  --ckpt-keep-steps 100000,200000,400000 \\\n  --ckpt-verify-step 0 \\\n  --ckpt-latest-freq 5000 \\\n  --no-fid-skip-eval \\\n  --ckpt-dir /kaggle/working/checkpoints/depth-shortcut-B-direct-boot-outputdistill-no-private-classcond-centered-logitnormal\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-29T18:15:54.585184Z","iopub.execute_input":"2026-04-29T18:15:54.585492Z","execution_failed":"2026-04-29T18:16:07.416Z"}},"outputs":[],"execution_count":null}]}