SprintML/MGI

1"""2Sample submission script for the MGI task.3 4End-to-end pipeline:5  1. Download the 900 reference images from the SprintML/MGI dataset repo6     on Hugging Face (data/img_000.png .. data/img_899.png).7  2. Download the RAR-XL generator weights (yucornetto/RAR) and the8     MaskGIT-VQ tokenizer weights.9  3. Build a valid 1800-slot submission.npz, following the same six-block10     transition convention in the task description:11       0000-0299 M->N   0300-0599 M->G   0600-0899 N->M12       0900-1199 N->G   1200-1499 G->M   1500-1799 G->N13  4. Submit it to the evaluation API.14 15build_submission() below is a placeholder that just returns random noise for16every slot. This is for familiarzing you with submission shape/format.17Please replace your own attack in its place.18"""19 20import os21import sys22import numpy as np23import requests24from pathlib import Path25from PIL import Image26from huggingface_hub import hf_hub_download, snapshot_download27 28BASE_DIR = Path(__file__).resolve().parent29 30# --- submission / API config -------------------------------------------------31BASE_URL    = "http://35.192.205.84"32API_KEY     = "YOUR_API_KEY_HERE"33TASK_ID     = "29-mgi"34OUTPUT_PATH = "submission.npz"35 36# --- submission format ---------------------------------------------------37# 1800 slots, six 300-image blocks for the six required misclassifications:38#   0000-0299 M->N   0300-0599 M->G   0600-0899 N->M39#   0900-1199 N->G   1200-1499 G->M   1500-1799 G->N40BASE_IMAGES  = 900   # underlying reference dataset (img_000.png .. img_899.png)41IMAGE_SIZE   = 25642TOTAL_IMAGES = 1800  # submission slots43EXPECTED_NAMES = tuple(f"{index:04d}" for index in range(TOTAL_IMAGES))44 45# --- Hugging Face sources -----------------------------------------------------46HF_DATASET_REPO   = "SprintML/MGI"47HF_DATA_SUBFOLDER = "data"48 49HF_RAR_REPO      = "yucornetto/RAR"       # RAR generator checkpoints (rar_xl.bin, ...)50HF_MASKGIT_REPO  = "fun-research/TiTok"   # MaskGIT-VQ tokenizer weight used by RAR51RAR_MODEL_SIZE   = "rar_xl"52 53MODEL_DIR = BASE_DIR / "model"54 55# RAR-XL architecture config: the hyperparameters RAR/demo_util.py needs to56# build the model class match those in the official RAR repo's57# configs/training/generator/rar.yaml for the XL size (see rar/README_RAR.md).58# Written locally only if not already present.59RAR_XL_CONFIG = """\60experiment:61    generator_checkpoint: ""62 63model:64    vq_model:65        codebook_size: 102466        token_size: 25667        num_latent_tokens: 25668        finetune_decoder: False69        pretrained_tokenizer_weight: ""70 71    generator:72        hidden_size: 128073        num_hidden_layers: 3274        num_attention_heads: 1675        intermediate_size: 512076        dropout: 0.177        attn_drop: 0.178        class_label_dropout: 0.179        image_seq_len: 25680        condition_num_classes: 100081        use_checkpoint: False82"""83 84 85def ensure_dataset() -> Path:86    """Download the 900 reference images from the HF dataset repo, if missing."""87    local_dir = snapshot_download(88        repo_id=HF_DATASET_REPO,89        repo_type="dataset",90        allow_patterns=[f"{HF_DATA_SUBFOLDER}/*.png"],91    )92    data_dir = Path(local_dir) / HF_DATA_SUBFOLDER93    print(f"Reference dataset ready: {data_dir}")94    return data_dir95 96 97def ensure_model_weights() -> tuple[Path, Path, Path]:98    """Download RAR-XL + MaskGIT-VQ weights and write a matching config, if missing."""99    MODEL_DIR.mkdir(parents=True, exist_ok=True)100 101    generator_ckpt = Path(102        hf_hub_download(repo_id=HF_RAR_REPO, filename=f"{RAR_MODEL_SIZE}.bin")103    )104    tokenizer_ckpt = Path(105        hf_hub_download(106            repo_id=HF_MASKGIT_REPO, filename="maskgit-vqgan-imagenet-f16-256.bin"107        )108    )109 110    config_path = MODEL_DIR / "rar.yaml"111    if not config_path.exists():112        config_path.write_text(RAR_XL_CONFIG)113 114    print(f"Model weights ready: generator={generator_ckpt}, tokenizer={tokenizer_ckpt}")115    return config_path, generator_ckpt, tokenizer_ckpt116 117 118def load_reference_images(data_dir: Path) -> np.ndarray:119    """Load the 900 reference dataset images as uint8 (BASE_IMAGES, 256, 256, 3)."""120    images = np.empty((BASE_IMAGES, IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8)121    for i in range(BASE_IMAGES):122        with Image.open(data_dir / f"img_{i:03d}.png") as img:123            images[i] = np.asarray(124                img.convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE), Image.BILINEAR),125                dtype=np.uint8,126            )127    return images128 129 130def build_submission(original: np.ndarray, seed: int = 0) -> np.ndarray:131    """132    Placeholder -- fill this in with your own attack.133 134    Returns random noise for every one of the 1800 slots, just to show the135    submission shape/format you need to produce. Scores 0 as-is.136    """137    rng = np.random.default_rng(seed)138    return rng.integers(139        0, 256, size=(TOTAL_IMAGES, IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8140    )141 142 143def make_submission_file(images: np.ndarray, output_path: str) -> None:144    assert images.shape == (TOTAL_IMAGES, IMAGE_SIZE, IMAGE_SIZE, 3), images.shape145    assert images.dtype == np.uint8, images.dtype146    names = np.array(EXPECTED_NAMES)147    np.savez_compressed(output_path, images=images, names=names)148    print(f"Saved submission -> {output_path}")149 150 151def die(msg: str) -> None:152    print(msg, file=sys.stderr)153    sys.exit(1)154 155 156def submit(file_path: str) -> None:157    if not os.path.isfile(file_path):158        die(f"File not found: {file_path}")159 160    try:161        with open(file_path, "rb") as f:162            files = {163                "file": (os.path.basename(file_path), f, "application/octet-stream"),164            }165            resp = requests.post(166                f"{BASE_URL}/submit/{TASK_ID}",167                headers={"X-API-Key": API_KEY},168                files=files,169            )170        try:171            body = resp.json()172        except Exception:173            body = {"raw_text": resp.text}174 175        if resp.status_code == 413:176            die("Upload rejected: file too large (HTTP 413). Reduce size and try again.")177 178        resp.raise_for_status()179 180        submission_id = body.get("submission_id")181        print("Successfully submitted.")182        print("Server response:", body)183        if submission_id:184            print(f"Submission ID: {submission_id}")185 186    except requests.exceptions.RequestException as e:187        detail = getattr(e, "response", None)188        print(f"Submission error: {e}")189        if detail is not None:190            try:191                print("Server response:", detail.json())192            except Exception:193                print("Server response (text):", detail.text)194        sys.exit(1)195 196 197if __name__ == "__main__":198    data_dir = ensure_dataset()199    ensure_model_weights()  # downloads RAR-XL + MaskGIT-VQ weights for your own attack200 201    original = load_reference_images(data_dir)202    submitted = build_submission(original)203    make_submission_file(submitted, OUTPUT_PATH)204    submit(OUTPUT_PATH)205