# GalaxyDiff Release

Self-contained **Measurement Alignment (MA)** inference package for conditional galaxy image generation.

Given 8 physical parameters, the package:

1. Samples a **64×64** image with a conditional diffusion model  
2. Upscales to **224×224** with SwinIR  
3. Measures the image with a **frozen** regression proxy  
4. Reports consistency metrics (MAE / RMSE / MAE\(_z\))

All paths resolve under this folder (`Paper_code/GalaxyDiff_release/`). No external data download is required once `weights/` is present (~2.5 GB).

---

## Pipeline

```
8-D conditions c
      │
      ▼
 CondNorm (z-score)
      │
      ▼
 MA diffusion  ──►  x₆₄  ──►  SwinIR SR  ──►  x₂₂₄
                                              │
                                              ▼
                                    frozen Reg proxy R(·)
                                              │
                                              ▼
                                         ĉ  vs  c
                                   MAE / RMSE / MAE_z
```

---

## Layout

```
GalaxyDiff_release/
├── generate_and_eval.py      # CLI entry
├── run.sh                    # one-command demo
├── demo_generate_and_eval.ipynb
├── README.md
├── ma/                       # release helpers
│   ├── pipeline.py           # end-to-end generate → SR → measure
│   ├── cond.py               # 8-D parse / CondNorm
│   ├── joint_ckpt.py         # load joint diffusion+SwinIR ckpt
│   ├── ddim.py               # fast sampling
│   ├── sr.py                 # SwinIR wrap
│   └── metrics.py            # MAE / RMSE / MAE_z
├── vendor/                   # model code only (no training)
│   ├── Daldiff_cond_norm/
│   ├── SwinIR-main/
│   └── Regression/
├── weights/                  # checkpoints (see MANIFEST.json)
│   ├── 00003000-joint.pt
│   ├── best_regression_head.pt
│   ├── cond_norm_stats.json
│   └── backbone/             # DINOv3 + LoRA / LN finetune
├── scripts/
│   └── assemble_bundle.sh    # pull weights/vendor from this repo (once)
└── runs/                     # example outputs
```

---

## Quick start

```bash
cd /home/zhangbi/LJM/diffusion/Paper_code/GalaxyDiff_release

# If weights/ is empty, assemble once from the parent repo:
bash scripts/assemble_bundle.sh

# Demo (GPU 0)
bash run.sh
```

Or call the CLI directly:

```bash
CUDA_VISIBLE_DEVICES=0 python generate_and_eval.py \
  --cond-values "0.0784,144.40,310.56,551.57,-0.3088,0.2504,5.1633,2.2935" \
  --n-samples 4 \
  --outdir runs/demo \
  --sampling-timesteps 50
```

Interactive walkthrough: open `demo_generate_and_eval.ipynb` from this directory.

---

## CLI options

| Flag | Default | Meaning |
|------|---------|---------|
| `--cond-values` | *(required)* | 8 comma-separated floats (order below) |
| `--n-samples` | `4` | number of stochastic samples |
| `--outdir` | `runs/demo` | output directory |
| `--device` | `cuda` | `cuda` or `cpu` |
| `--rank` | `0` | GPU index for `DataParallel` |
| `--base-seed` | `42` | RNG seed for sampling |
| `--batch-size` | `4` | generation batch size |
| `--sampling-timesteps` | `50` | DDIM steps; `0` or `1000` → full 1000-step DDPM |
| `--mc-samples` | `10` | MC Dropout passes in the proxy |
| `--skip-generate` | off | reuse existing LR images in `outdir` |
| `--skip-sr` | off | skip SwinIR (LR-only path) |

`run.sh` env overrides: `COND_VALUES`, `N_SAMPLES`, `OUTDIR`, `SAMPLING_TIMESTEPS`, `MC_SAMPLES`, `CUDA_VISIBLE_DEVICES`.

---

## Input conditions (fixed order)

| # | Name | Meaning |
|---|------|---------|
| 1 | `redshift` | spectroscopic / photometric redshift |
| 2 | `flux_g` | g-band flux |
| 3 | `flux_r` | r-band flux |
| 4 | `flux_z` | z-band flux |
| 5 | `shape_e1` | ellipticity \(e_1\) |
| 6 | `shape_e2` | ellipticity \(e_2\) |
| 7 | `shape_r` | effective radius \(r_e\) |
| 8 | `sersic` | Sérsic index |

Values are **physical units** (not z-scored). CondNorm stats live in `weights/cond_norm_stats.json`.

---

## Outputs (`--outdir`)

| Path | Content |
|------|---------|
| `lr/output/*.png` | generated 64×64 images |
| `sr/sr/*.png` | 224×224 super-resolved images |
| `param_comparison.csv` | input \(c\) vs measured \(\hat{c}\) |
| `summary.json` / `metrics.json` | MAE, RMSE, MAE\(_z\) (when \(N\ge 2\)) |
| `input_cond.json` | echo of the requested conditions |

---

## Weights

See `weights/MANIFEST.json`:

| File | Role |
|------|------|
| `00003000-joint.pt` | joint MA checkpoint (diffusion + SwinIR, step 3000) |
| `best_regression_head.pt` | frozen measurement proxy |
| `cond_norm_stats.json` | condition z-score mean / std |
| `backbone/dinov3-vitl16-…` | DINOv3 ViT-L/16 |
| `backbone/finetune_2addLN/` | LoRA + LN adapters for the proxy |

Approximate disk: **~2.5 GB** under `weights/`.

---

## Assemble from this repository

Only needed if `weights/` or `vendor/` are missing (e.g. after a fresh clone without the bundle):

```bash
cd /home/zhangbi/LJM/diffusion/Paper_code/GalaxyDiff_release
bash scripts/assemble_bundle.sh
```

The script copies / hardlinks diffusion, SwinIR, and regression code from `diffusion_model/`, and checkpoints from related `Paper_code/experiment*` paths.

---

## Dependencies

```
torch
torchvision
transformers
timm
Pillow
numpy
pandas
matplotlib
tqdm
```

CUDA GPU recommended. Full DDPM (`--sampling-timesteps 0`) is much slower than DDIM 50.

---

## Notes

- This package is **inference-only**; training lives in `diffusin_model_fits/` and `diffusion_model/`.
- Folder was formerly named `model_release`; paths in older docs may still say that — use `GalaxyDiff_release` instead.
- Prefer `--sampling-timesteps 50` for demos; use `1000` / `0` only when comparing to paper-quality DDPM sampling.