{"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":"none","dataSources":[{"sourceId":6799,"databundleVersionId":4225553,"isSourceIdPinned":false,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport re\nfrom pathlib import Path\n\nINPUT_ROOT = Path(\"/kaggle/input/imagenet-object-localization-challenge/ILSVRC/Data/CLS-LOC\")\nTRAIN_SRC = INPUT_ROOT / \"train\"\nVAL_SRC   = INPUT_ROOT / \"val\"\nMAP_FILE  = Path(\"/kaggle/input/imagenet-object-localization-challenge/LOC_val_solution.csv\")\n\nWORK_ROOT = Path(\"/kaggle/working/imagenet\")\nWORK_TRAIN = WORK_ROOT / \"train\"\nWORK_VAL   = WORK_ROOT / \"val\"\n\nassert TRAIN_SRC.exists(), f\"Missing: {TRAIN_SRC}\"\nassert VAL_SRC.exists(), f\"Missing: {VAL_SRC}\"\nassert MAP_FILE.exists(), f\"Missing: {MAP_FILE}\"\n\nWORK_ROOT.mkdir(parents=True, exist_ok=True)\n\ndef symlink_if_needed(src: Path, dst: Path):\n    dst.parent.mkdir(parents=True, exist_ok=True)\n    if dst.exists() or dst.is_symlink():\n        return\n    os.symlink(src.as_posix(), dst.as_posix())\n\ndef parse_loc_val_solution(map_path: Path):\n    \"\"\"\n    Trả về dict: {ImageId(without .JPEG): wnid}\n    Hỗ trợ:\n      - CSV chuẩn: ImageId,PredictionString\n      - Kiểu 2 dòng: ImageId \\n \"wnid x1 y1 x2 y2 ...\"\n    \"\"\"\n    mapping = {}\n\n    # Đọc vài dòng đầu để đoán format\n    head_lines = []\n    with map_path.open(\"r\", encoding=\"utf-8\", errors=\"ignore\") as f:\n        for _ in range(5):\n            line = f.readline()\n            if not line:\n                break\n            head_lines.append(line.strip())\n\n    is_csv_header = any(\"ImageId\" in l for l in head_lines) and any(\",\" in l for l in head_lines)\n\n    if is_csv_header:\n        # CSV dạng Kaggle phổ biến\n        import csv\n        with map_path.open(\"r\", encoding=\"utf-8\", errors=\"ignore\", newline=\"\") as f:\n            reader = csv.DictReader(f)\n            if \"ImageId\" not in reader.fieldnames:\n                raise ValueError(f\"CSV header không có ImageId: {reader.fieldnames}\")\n            # Thường là PredictionString\n            pred_col = \"PredictionString\" if \"PredictionString\" in reader.fieldnames else None\n            if pred_col is None:\n                # fallback: lấy cột thứ 2\n                cols = [c for c in reader.fieldnames if c != \"ImageId\"]\n                if not cols:\n                    raise ValueError(f\"Không tìm thấy cột prediction trong CSV: {reader.fieldnames}\")\n                pred_col = cols[0]\n\n            for row in reader:\n                imgid = (row.get(\"ImageId\") or \"\").strip()\n                pred  = (row.get(pred_col) or \"\").strip()\n                if not imgid or not pred:\n                    continue\n                # PredictionString có thể chứa nhiều object: wnid x1 y1 x2 y2 wnid x1 y1 x2 y2 ...\n                wnid = pred.split()[0]\n                if re.fullmatch(r\"n\\d{8}\", wnid):\n                    mapping[imgid] = wnid\n        return mapping\n\n    # Fallback: parse kiểu “2 dòng” như bạn paste\n    current_imgid = None\n    with map_path.open(\"r\", encoding=\"utf-8\", errors=\"ignore\") as f:\n        for raw in f:\n            line = raw.strip()\n            if not line:\n                continue\n\n            # Nếu là ImageId (không có khoảng trắng), ví dụ ILSVRC2012_val_00048981\n            if re.fullmatch(r\"ILSVRC2012_val_\\d{8}\", line):\n                current_imgid = line\n                continue\n\n            # Nếu là dòng prediction bắt đầu bằng wnid\n            if current_imgid is not None:\n                first = line.split()[0]\n                if re.fullmatch(r\"n\\d{8}\", first):\n                    mapping[current_imgid] = first\n                    current_imgid = None\n\n    if not mapping:\n        raise ValueError(\"Không parse được LOC_val_solution.csv. Hãy mở vài dòng đầu của file để kiểm tra format.\")\n    return mapping\n\n# 1) Symlink train\nif not WORK_TRAIN.exists():\n    os.symlink(TRAIN_SRC.as_posix(), WORK_TRAIN.as_posix())\nprint(\"train symlink:\", WORK_TRAIN, \"->\", os.readlink(WORK_TRAIN))\n\n# 2) Parse map\nimgid2wnid = parse_loc_val_solution(MAP_FILE)\nprint(\"Parsed mappings:\", len(imgid2wnid))\n\n# 3) Tạo trước 1000 folder wnid theo train để đảm bảo class order nhất quán\nWORK_VAL.mkdir(parents=True, exist_ok=True)\ntrain_wnids = sorted([p.name for p in TRAIN_SRC.iterdir() if p.is_dir() and re.fullmatch(r\"n\\d{8}\", p.name)])\nprint(\"Train wnids:\", len(train_wnids))\nfor wnid in train_wnids:\n    (WORK_VAL / wnid).mkdir(parents=True, exist_ok=True)\n\n# 4) Symlink val ảnh vào đúng wnid folder\nmissing_img = 0\nlinked = 0\n\nfor imgid, wnid in imgid2wnid.items():\n    src = VAL_SRC / f\"{imgid}.JPEG\"\n    if not src.exists():\n        # đôi khi extension khác case\n        src = VAL_SRC / f\"{imgid}.jpeg\"\n    if not src.exists():\n        missing_img += 1\n        continue\n\n    dst = WORK_VAL / wnid / src.name\n    if not dst.exists():\n        os.symlink(src.as_posix(), dst.as_posix())\n        linked += 1\n\nprint(\"Linked val images:\", linked)\nprint(\"Missing val images:\", missing_img)\n\n# 5) Sanity check nhanh\nval_count = sum(1 for _ in WORK_VAL.rglob(\"*.JPEG\")) + sum(1 for _ in WORK_VAL.rglob(\"*.jpeg\"))\nprint(\"WORK_VAL total images:\", val_count)\nprint(\"WORK_ROOT:\", WORK_ROOT)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T17:27:45.990345Z","iopub.execute_input":"2025-12-25T17:27:45.990792Z","iopub.status.idle":"2025-12-25T17:30:01.705407Z","shell.execute_reply.started":"2025-12-25T17:27:45.990751Z","shell.execute_reply":"2025-12-25T17:30:01.70414Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install -q tfds apache_beam mlcroissant\n# !curl -LsSf https://astral.sh/uv/install.sh | sh\n# os.environ[\"PATH\"] += \":/root/.local/bin\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T17:33:55.836124Z","iopub.execute_input":"2025-12-25T17:33:55.836582Z","iopub.status.idle":"2025-12-25T17:34:55.330212Z","shell.execute_reply.started":"2025-12-25T17:33:55.836554Z","shell.execute_reply":"2025-12-25T17:34:55.328664Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !git clone https://github.com/Gsunshine/meanflow\n# %cd meanflow","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T17:34:55.332384Z","iopub.execute_input":"2025-12-25T17:34:55.332726Z","iopub.status.idle":"2025-12-25T17:34:56.035473Z","shell.execute_reply.started":"2025-12-25T17:34:55.332695Z","shell.execute_reply":"2025-12-25T17:34:56.034336Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%bash\n# uv python install 3.10\n# uv python pin 3.10\n# uv venv .venv --python 3.10\n# uv init\n# uv add pip\n# uv run bash","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T17:34:56.036792Z","iopub.execute_input":"2025-12-25T17:34:56.037094Z","iopub.status.idle":"2025-12-25T17:34:59.637443Z","shell.execute_reply.started":"2025-12-25T17:34:56.037067Z","shell.execute_reply":"2025-12-25T17:34:59.636337Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# with open(\"/kaggle/working/meanflow/scripts/install.sh\", \"w\") as file:\n#     file.write(r\"\"\"uv run pip install jax[tpu]==0.4.27 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n# uv run pip install jaxlib==0.4.27 \"flax>=0.8\"\n# uv run pip install pillow clu tensorflow==2.15.0 \"keras<3\" \"torch<=2.4\" torchvision tensorflow_datasets matplotlib==3.9.2\n# uv run pip install orbax-checkpoint==0.4.4 ml-dtypes==0.5.0 tensorstore==0.1.67\n# uv run pip install diffusers dm-tree cached_property\"\"\")\n# # with open(\"/kaggle/working/meanflow/scripts/install.sh\", \"w\") as file:\n# #     file.write(r\"\"\"uv run pip install jax[tpu]==0.4.13 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n# # uv run pip install jaxlib==0.4.13 \"flax>=0.7.2\"\n# # uv run pip install pillow clu tensorflow==2.13.1 \"keras<3\" \"torch<=2.4\" torchvision tensorflow_datasets matplotlib==3.7.5\n# # uv run pip install orbax-checkpoint==0.4.4 ml-dtypes==0.5.0 tensorstore==0.1.67\n# # uv run pip install diffusers dm-tree cached_property\"\"\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T17:40:34.964861Z","iopub.execute_input":"2025-12-25T17:40:34.965313Z","iopub.status.idle":"2025-12-25T17:40:35.03541Z","shell.execute_reply.started":"2025-12-25T17:40:34.965276Z","shell.execute_reply":"2025-12-25T17:40:35.034287Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !bash \"/kaggle/working/meanflow/scripts/install.sh\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T17:38:42.188758Z","iopub.execute_input":"2025-12-25T17:38:42.189168Z","iopub.status.idle":"2025-12-25T17:40:34.961035Z","shell.execute_reply.started":"2025-12-25T17:38:42.189142Z","shell.execute_reply":"2025-12-25T17:40:34.95917Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# file_path = \"/kaggle/working/meanflow/scripts/prepare_data.sh\"\n\n# with open(file_path,\"w\") as file:\n#     file.write(\"\"\"#!/bin/bash\n\n# # Configuration for data preparation\n# export IMAGENET_ROOT=\"YOUR_IMAGENET_ROOT\"\n# export OUTPUT_DIR=\"YOUR_OUTPUT_DIR\"\n# export LOG_DIR=\"YOUR_LOG_DIR\"\n\n# # Validate required environment variables\n# if [ \"$IMAGENET_ROOT\" = \"YOUR_IMAGENET_ROOT\" ] || [ \"$OUTPUT_DIR\" = \"YOUR_OUTPUT_DIR\" ] || [ \"$LOG_DIR\" = \"YOUR_LOG_DIR\" ]; then\n#     echo \"ERROR: Please update the environment variables at the top of this script:\"\n#     echo \"  - IMAGENET_ROOT: Path to your ImageNet dataset\"\n#     echo \"  - OUTPUT_DIR: Path where to save the processed data\"\n#     echo \"  - LOG_DIR: Path where to save logs\"\n#     exit 1\n# fi\n\n# export BATCH_SIZE=128\n# export VAE_TYPE=\"mse\"\n\n# export now=`date '+%Y%m%d_%H%M%S'`\n# export salt=`head /dev/urandom | tr -dc a-z0-9 | head -c6`\n# export JOBNAME=prepare_data_${now}_${salt}_$1\n# export LOG_DIR=$LOG_DIR/$USER/$JOBNAME\n\n# sudo mkdir -p ${LOG_DIR}\n# sudo chmod 777 -R ${LOG_DIR}\n\n# # Image size configuration (common sizes: 256, 512, 1024)\n# # Corresponding latent sizes will be: 32x32, 64x64, 128x128\n# IMAGE_SIZE=${IMAGE_SIZE:-256}  # Can be overridden via environment variable\n\n# # Computation flags (can be overridden via environment variables)\n# COMPUTE_LATENT=${COMPUTE_LATENT:-True}  # Whether to compute latent dataset\n# COMPUTE_FID=${COMPUTE_FID:-False}       # Whether to compute FID statistics\n\n# # Calculate latent size for display\n# LATENT_SIZE=$((IMAGE_SIZE / 8))\n\n# echo \"==============================================\"\n# echo \"Data Preparation Configuration\"\n# echo \"==============================================\"\n# echo \"ImageNet Root: $IMAGENET_ROOT\"\n# echo \"Output Dir: $OUTPUT_DIR\"\n# echo \"Batch Size: $BATCH_SIZE\"\n# echo \"VAE Type: $VAE_TYPE\"\n# echo \"Image Size: $IMAGE_SIZE -> Latent Size: ${LATENT_SIZE}x${LATENT_SIZE}\"\n# echo \"Compute Latent: $COMPUTE_LATENT\"\n# echo \"Compute FID: $COMPUTE_FID\"\n# if [ \"$COMPUTE_FID\" = \"True\" ]; then\n#     echo \"FID: Using ALL training samples\"\n# fi\n# echo \"==============================================\"\n\n# uv run prepare_dataset.py \\\n#     --imagenet_root=\\\"$IMAGENET_ROOT\\\" \\\n#     --output_dir=\\\"$OUTPUT_DIR\\\" \\\n#     --batch_size=$BATCH_SIZE \\\n#     --vae_type=\\\"$VAE_TYPE\\\" \\\n#     --image_size=$IMAGE_SIZE \\\n#     --compute_latent=$COMPUTE_LATENT \\\n#     --compute_fid=$COMPUTE_FID \\\n#     --overwrite=False \\\n#     2>&1 | tee -a $LOG_DIR/output.log\n\n# echo \"==============================================\"\n# echo \"Data preparation completed!\"\n# echo \"Check logs at: $LOG_DIR/output.log\"\n# if [ \"$COMPUTE_LATENT\" = \"True\" ]; then\n#     echo \"Latent dataset saved to: $OUTPUT_DIR\"\n# fi\n# if [ \"$COMPUTE_FID\" = \"True\" ]; then\n#     echo \"FID stats saved to: $OUTPUT_DIR/imagenet_${IMAGE_SIZE}_fid_stats.npz\"\n# fi\n# echo \"==============================================\" \"\"\")\n# !chmod +x {file_path}\n# !bash {file_path}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T17:34:59.64673Z","iopub.execute_input":"2025-12-25T17:34:59.64711Z","iopub.status.idle":"2025-12-25T17:34:59.915335Z","shell.execute_reply.started":"2025-12-25T17:34:59.647059Z","shell.execute_reply":"2025-12-25T17:34:59.91413Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import os\n# from pathlib import Path\n\n# # ==== chỉnh 3 biến này cho đúng notebook của bạn ====\n# IMAGENET_ROOT = \"/kaggle/working/imagenet\"          # root chứa train/ và (val/ nếu cần)\n# OUTPUT_DIR    = \"/kaggle/working/meanflow_data\"     # nơi lưu latent .pt\n# BASE_LOG_DIR  = \"/kaggle/working/logs\"              # nơi lưu log\n# MEANFLOW_DIR  = \"/kaggle/working/meanflow\"          # repo meanflow của bạn (để chạy prepare_dataset.py)\n\n# # sanity check nhanh\n# assert Path(IMAGENET_ROOT).exists(), f\"IMAGENET_ROOT not found: {IMAGENET_ROOT}\"\n# assert Path(MEANFLOW_DIR).exists(), f\"MEANFLOW_DIR not found: {MEANFLOW_DIR}\"\n\n# scripts_dir = Path(MEANFLOW_DIR) / \"scripts\"\n# scripts_dir.mkdir(parents=True, exist_ok=True)\n\n# file_path = scripts_dir / \"prepare_data.sh\"\n\n# script = f\"\"\"#!/bin/bash\n# set -euo pipefail\n\n# export JAX_PROCESS_COUNT=1\n# export JAX_PROCESS_INDEX=0\n# export JAX_COORDINATOR_ADDRESS=127.0.0.1:1234\n\n# export IMAGENET_ROOT=\"{IMAGENET_ROOT}\"\n# export OUTPUT_DIR=\"{OUTPUT_DIR}\"\n# export BASE_LOG_DIR=\"{BASE_LOG_DIR}\"\n\n# export BATCH_SIZE=\"${{BATCH_SIZE:-128}}\"\n# export VAE_TYPE=\"${{VAE_TYPE:-mse}}\"\n\n# export now=$(date '+%Y%m%d_%H%M%S')\n# export salt=$(head /dev/urandom | tr -dc a-z0-9 | head -c6)\n# export JOBNAME=prepare_data_${{now}}_${{salt}}\n\n# export LOG_DIR=\"${{BASE_LOG_DIR}}/${{JOBNAME}}\"\n# mkdir -p \"${{LOG_DIR}}\"\n# chmod 777 -R \"${{LOG_DIR}}\" || true\n\n# # Image size configuration (common sizes: 256, 512, 1024)\n# IMAGE_SIZE=\"${{IMAGE_SIZE:-256}}\"\n\n# # Computation flags (override by env)\n# COMPUTE_LATENT=\"${{COMPUTE_LATENT:-True}}\"\n# COMPUTE_FID=\"${{COMPUTE_FID:-False}}\"\n\n# LATENT_SIZE=$((IMAGE_SIZE / 8))\n\n# echo \"==============================================\"\n# echo \"Data Preparation Configuration\"\n# echo \"==============================================\"\n# echo \"ImageNet Root: $IMAGENET_ROOT\"\n# echo \"Output Dir: $OUTPUT_DIR\"\n# echo \"Log Dir: $LOG_DIR\"\n# echo \"Batch Size: $BATCH_SIZE\"\n# echo \"VAE Type: $VAE_TYPE\"\n# echo \"Image Size: $IMAGE_SIZE -> Latent Size: ${{LATENT_SIZE}}x${{LATENT_SIZE}}\"\n# echo \"Compute Latent: $COMPUTE_LATENT\"\n# echo \"Compute FID: $COMPUTE_FID\"\n# echo \"==============================================\"\n\n# mkdir -p \"$OUTPUT_DIR\"\n\n# cd \"{MEANFLOW_DIR}\"\n\n# # Prefer uv if available, else fallback to python\n# RUNNER=\"python\"\n# if command -v uv >/dev/null 2>&1; then\n#   RUNNER=\"uv run\"\n# fi\n\n# $RUNNER prepare_dataset.py \\\\\n#   --imagenet_root=\"$IMAGENET_ROOT\" \\\\\n#   --output_dir=\"$OUTPUT_DIR\" \\\\\n#   --batch_size=\"$BATCH_SIZE\" \\\\\n#   --vae_type=\"$VAE_TYPE\" \\\\\n#   --image_size=\"$IMAGE_SIZE\" \\\\\n#   --compute_latent=\"$COMPUTE_LATENT\" \\\\\n#   --compute_fid=\"$COMPUTE_FID\" \\\\\n#   --overwrite=False \\\\\n#   2>&1 | tee -a \"$LOG_DIR/output.log\"\n\n# echo \"==============================================\"\n# echo \"Data preparation completed!\"\n# echo \"Check logs at: $LOG_DIR/output.log\"\n# echo \"==============================================\"\n# \"\"\"\n\n# file_path.write_text(script)\n# os.chmod(file_path, 0o755)\n\n# print(\"Wrote:\", file_path)\n# print(\"Running...\")\n\n# # Bạn có thể override env ở đây nếu muốn:\n# # os.environ[\"BATCH_SIZE\"] = \"64\"\n# # os.environ[\"IMAGE_SIZE\"] = \"256\"\n# # os.environ[\"COMPUTE_LATENT\"] = \"True\"\n# # os.environ[\"COMPUTE_FID\"] = \"False\"\n\n# os.system(f\"bash {file_path}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T17:43:17.379098Z","iopub.execute_input":"2025-12-25T17:43:17.379553Z","iopub.status.idle":"2025-12-25T17:43:18.454295Z","shell.execute_reply.started":"2025-12-25T17:43:17.379513Z","shell.execute_reply":"2025-12-25T17:43:18.452829Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%writefile pyproject.toml\n# [project]\n# name = \"shortcut-models\"\n# version = \"0.1.0\"\n# description = \"Add your description here\"\n# readme = \"README.md\"\n# requires-python = \"==3.11.6\"\n# dependencies = [\n#     \"absl-py>=2.3.1\",\n#     \"chex>=0.1.86\",\n#     \"cython<3\",\n#     \"diffusers>=0.35.2\",\n#     \"distrax==0.1.4\",\n#     \"einops>=0.8.1\",\n#     \"fabric>=3.2.2\",\n#     \"flax>=0.8.3\",\n#     \"imageio>=2.37.0\",\n#     \"jax[tpu]==0.4.27\",\n#     \"jaxtyping>=0.3.3\",\n#     \"libtmux>=0.46.2\",\n#     \"matplotlib>=3.10.7\",\n#     \"ml-collections>=1.1.0\",\n#     \"moviepy>=2.2.1\",\n#     \"numba>=0.62.1\",\n#     \"numpy>=1.26.4\",\n#     \"opensimplex>=0.4.5.1\",\n#     \"opt-einsum>=3.4.0\",\n#     \"optax<=0.2.4\",\n#     \"orbax==0.1.9\",\n#     \"plotly>=6.3.1\",\n#     \"protobuf<=3.20.3\",\n#     \"pygame>=2.6.1\",\n#     \"scipy>1.12.0\",\n#     \"tabulate>=0.9.0\",\n#     \"tensorflow-cpu>=2.16.0\",\n#     \"tensorflow-datasets>=4.9.9\",\n#     \"tensorflow-probability==0.22.0\",\n#     \"termcolor>=3.1.0\",\n#     \"threadpoolctl==3.1.0\",\n#     \"typeguard>=4.0.0\",\n#     \"wandb>=0.19.11\",\n#     \"wheel>=0.45.1\",\n# ]","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}