Upload 4 files

run_demo.py

@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+"""
+run_demo.py — Self-contained dummy demo of upload_to_hub.py
+============================================================
+Builds a fake HF deployment package WITHOUT requiring torch or a real
+checkpoint, so you can see exactly what files get uploaded.
+
+This demo:
+  1. Creates a dummy checkpoint, args.json, label stats files
+  2. Patches torch import to a stub so upload_to_hub.py can run
+  3. Calls package_model() in dry-run mode
+  4. Lists every file in the package with its purpose
+
+Run:
+  python run_demo.py
+"""
+from __future__ import annotations
+
+import json
+import shutil
+import sys
+import types
+from pathlib import Path
+
+import numpy as np
+
+
+# ── §1  Build a torch stub (so upload_to_hub.py can be imported) ───────────
+
+class _TorchStub:
+    class Tensor:
+        def __init__(self, data):
+            self._d = np.asarray(data)
+            self.shape = self._d.shape
+        def numel(self):  return int(np.prod(self.shape))
+        def clone(self):  return self
+        def contiguous(self): return self
+        @property
+        def dtype(self):  return _DType()
+    @staticmethod
+    def load(path, **kw):
+        # Simulate loading our dummy checkpoint
+        return _DUMMY_CKPT
+    @staticmethod
+    def save(obj, path):
+        # Mimic torch.save — for the .bin fallback path
+        with open(path, "wb") as f:
+            f.write(b"DUMMY_TORCH_BIN")
+
+class _DType:
+    @property
+    def is_floating_point(self): return True
+
+
+# Mock checkpoint structure that mirrors a real DDPM checkpoint
+_DUMMY_CKPT = {
+    "model_state_dict": {
+        "unet.conv.weight":     _TorchStub.Tensor(np.zeros((64, 1, 3, 3), dtype=np.float32)),
+        "unet.conv.bias":       _TorchStub.Tensor(np.zeros(64, dtype=np.float32)),
+        "unet.label_emb.weight":_TorchStub.Tensor(np.zeros((64, 2), dtype=np.float32)),
+        "unet.label_emb.bias":  _TorchStub.Tensor(np.zeros(64, dtype=np.float32)),
+        "unet.out.weight":      _TorchStub.Tensor(np.zeros((1, 64, 1, 1), dtype=np.float32)),
+        "unet.out.bias":        _TorchStub.Tensor(np.zeros(1, dtype=np.float32)),
+    },
+    "ema_shadow": {
+        "unet.conv.weight":     _TorchStub.Tensor(np.ones((64, 1, 3, 3), dtype=np.float32)*0.01),
+        "unet.conv.bias":       _TorchStub.Tensor(np.zeros(64, dtype=np.float32)),
+        "unet.label_emb.weight":_TorchStub.Tensor(np.zeros((64, 2), dtype=np.float32)),
+        "unet.label_emb.bias":  _TorchStub.Tensor(np.zeros(64, dtype=np.float32)),
+        "unet.out.weight":      _TorchStub.Tensor(np.zeros((1, 64, 1, 1), dtype=np.float32)),
+        "unet.out.bias":        _TorchStub.Tensor(np.zeros(1, dtype=np.float32)),
+    },
+    "epoch": 100,
+}
+
+# Stub safetensors too (writes a fake binary blob)
+class _SafetensorsStub:
+    @staticmethod
+    def save_file(state_dict, path):
+        # Just write a fake header so file exists with realistic size
+        # In reality safetensors writes a JSON header + binary tensor data
+        total_bytes = sum(t.numel() * 4 for t in state_dict.values())
+        with open(path, "wb") as f:
+            f.write(b"\x00" * total_bytes)
+
+
+# ── §2  Set up the dummy project ───────────────────────────────────────────
+
+DEMO_ROOT = Path("/tmp/ddpm_hf_demo")
+PROJECT   = DEMO_ROOT / "project"
+EXPORT    = DEMO_ROOT / "hf_export"
+
+if DEMO_ROOT.exists():
+    shutil.rmtree(DEMO_ROOT)
+PROJECT.mkdir(parents=True)
+(PROJECT / "checkpoints").mkdir()
+
+# Minimal source files (will be copied into the HF package)
+(PROJECT / "diffusion_conditional.py").write_text(
+    '"""Stub: our DDPM forward/reverse process implementation."""\n'
+    'import torch.nn as nn\n'
+    'class GaussianDiffusion(nn.Module): ...\n'
+    'class ConditionalDiffusionModel(nn.Module): ...\n'
+)
+(PROJECT / "unet_conditional.py").write_text(
+    '"""Stub: our conditional U-Net architecture."""\n'
+    'import torch.nn as nn\n'
+    'class ConditionalUNet(nn.Module): ...\n'
+)
+
+# Fake checkpoint (file content doesn't matter — torch.load is stubbed)
+(PROJECT / "checkpoints/best_model.pt").write_bytes(b"DUMMY_CKPT")
+
+# Training config
+(PROJECT / "args.json").write_text(json.dumps({
+    "image_size": 256, "label_dim": 2,
+    "base_channels": 64, "channel_multipliers": [1, 2, 4, 8],
+    "attention_levels": [2, 3], "dropout": 0.1,
+    "timesteps": 1500, "beta_start": 1e-4, "beta_end": 0.02,
+    "schedule_type": "linear", "ddim_steps": 50,
+    "epochs": 100, "batch_size": 8, "lr": 2e-4,
+    "ema_decay": 0.9999, "seed": 42,
+}, indent=2))
+
+# Training labels (for label_mu / label_std extraction)
+labels = np.random.uniform([0.1, 0.6], [0.5, 1.0], (50, 2)).astype(np.float32)
+np.save(PROJECT / "train_labels_LH_2.npy", labels)
+
+
+# ── §3  Inject stubs into sys.modules and import upload_to_hub ─────────────
+
+sys.modules["torch"] = _TorchStub()
+sys.modules["safetensors"] = types.ModuleType("safetensors")
+sys.modules["safetensors.torch"] = _SafetensorsStub()
+
+# Also stub huggingface_hub so we don't hit the network
+class _HfStub:
+    HfApi = type("HfApi", (), {
+        "create_repo": lambda *a, **kw: None,
+        "upload_folder": lambda *a, **kw: None,
+    })
+    login = lambda *a, **kw: None
+sys.modules["huggingface_hub"] = _HfStub()
+
+sys.path.insert(0, str(Path(__file__).parent))
+import upload_to_hub
+
+
+# ── §4  Run package_model() in dry-run mode ────────────────────────────────
+
+class FakeArgs:
+    checkpoint    = str(PROJECT / "checkpoints/best_model.pt")
+    training_args = str(PROJECT / "args.json")
+    data_dir      = str(PROJECT)
+    export_dir    = str(EXPORT)
+    no_ema        = False
+    repo_id       = "demo-user/camels-ddpm-omega-sigma8"
+
+print("="*65)
+print("  DDPM -> Hugging Face Hub Packager  (DUMMY DEMO)")
+print("="*65)
+folder = upload_to_hub.package_model(FakeArgs())
+
+
+# ── §5  Verify the result ──────────────────────────────────────────────────
+
+print("\n" + "="*65)
+print("  Package verification")
+print("="*65)
+
+config = json.loads((folder / "config.json").read_text())
+print("\nconfig.json contents:")
+print(json.dumps(config, indent=2))
+
+print(f"\nREADME.md preview (first 50 lines):")
+print("-"*65)
+print("\n".join((folder / "README.md").read_text().splitlines()[:50]))
+print("...")
+print("-"*65)
+
+print(f"\nFile listing of {folder}:")
+files = sorted(folder.iterdir())
+print(f"\n{'File':<32} {'Size':>10}  Purpose")
+print("-"*75)
+purposes = {
+    "config.json":              "Architecture hyperparameters (hub-readable)",
+    "model.safetensors":        "Model weights (EMA preferred)",
+    "pytorch_model.bin":        "Model weights (fallback if no safetensors)",
+    "README.md":                "Model card with YAML metadata + usage docs",
+    "modeling_ddpm_camels.py":  "Self-contained loader for `from_pretrained`",
+    "diffusion_conditional.py": "Project file: forward/reverse DDPM process",
+    "unet_conditional.py":      "Project file: U-Net architecture",
+    "inference_example.py":     "Standalone demo script for users",
+    "requirements.txt":         "Pinned Python dependencies",
+    ".gitattributes":           "Git LFS configuration for large files",
+}
+for f in files:
+    sz = f.stat().st_size
+    sz_str = f"{sz/1e6:.1f}M" if sz > 1e6 else f"{sz/1e3:.1f}K" if sz > 1e3 else f"{sz}B"
+    purpose = purposes.get(f.name, "")
+    print(f"  {f.name:<30} {sz_str:>10}  {purpose}")
+
+print(f"\nDemo complete -> {folder}")
+print(f"In a real run, the next step is:")
+print(f"  python upload_to_hub.py --checkpoint best_model.pt \\")
+print(f"      --training_args args.json \\")
+print(f"      --repo_id YOUR_USERNAME/camels-ddpm \\")
+print(f"      --private")

sample_config.json

@@ -0,0 +1,45 @@
+{
+  "model_type": "conditional_ddpm_camels",
+  "in_channels": 1,
+  "out_channels": 1,
+  "image_size": 256,
+  "label_dim": 2,
+  "label_names": [
+    "Omega_m",
+    "sigma_8"
+  ],
+  "base_channels": 64,
+  "channel_multipliers": [
+    1,
+    2,
+    4,
+    8
+  ],
+  "attention_levels": [
+    2,
+    3
+  ],
+  "dropout": 0.1,
+  "timesteps": 1500,
+  "beta_start": 0.0001,
+  "beta_end": 0.02,
+  "schedule_type": "linear",
+  "ddim_steps_default": 50,
+  "framework": "pytorch",
+  "library_name": "pytorch",
+  "training_meta": {
+    "epochs": 100,
+    "batch_size": 8,
+    "lr": 0.0002,
+    "ema_decay": 0.9999,
+    "seed": 42
+  },
+  "label_mu": [
+    0.3308129608631134,
+    0.7831979990005493
+  ],
+  "label_std": [
+    0.1140434592962265,
+    0.12279357761144638
+  ]
+}

sample_inference_example.py

@@ -0,0 +1,62 @@
+"""
+inference_example.py
+====================
+Standalone script demonstrating how to use the deployed DDPM model.
+After downloading from the Hub, run:
+    python inference_example.py
+"""
+import json
+import sys
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+# Ensure local imports resolve
+sys.path.insert(0, str(Path(__file__).parent))
+
+from modeling_ddpm_camels import load_pretrained, generate
+
+# ── Configuration ──────────────────────────────────────────────────────────
+MODEL_DIR = Path(__file__).parent
+DEVICE    = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ── Load ───────────────────────────────────────────────────────────────────
+print(f"Loading model from {MODEL_DIR} on {DEVICE} ...")
+model, config = load_pretrained(MODEL_DIR, device=DEVICE)
+print(f"  Image size: {config[\"image_size\"]}")
+print(f"  Label dim:  {config[\"label_dim\"]} ({config[\"label_names\"]})")
+
+# ── Generate at 4 cosmologies ──────────────────────────────────────────────
+raw_labels = torch.tensor([
+    [0.20, 0.95],
+    [0.30, 0.80],
+    [0.40, 0.70],
+    [0.50, 0.65],
+], dtype=torch.float32)
+
+if config["label_dim"] > 2:
+    # Pad with fiducial astrophysics (label_mu values of those dims)
+    pad = torch.tensor(config["label_mu"][2:], dtype=torch.float32).unsqueeze(0)
+    raw_labels = torch.cat([raw_labels, pad.expand(4, -1)], dim=1)
+
+print(f"\nGenerating samples ...")
+with torch.no_grad():
+    out = generate(model, config, raw_labels, device=DEVICE, ddim_steps=50)
+
+# Map [-1, 1] -> [0, 1] for visualisation
+imgs = ((out.cpu().numpy() + 1) / 2).clip(0, 1)[:, 0]
+
+# ── Display ────────────────────────────────────────────────────────────────
+fig, axes = plt.subplots(1, len(imgs), figsize=(3 * len(imgs), 3.5))
+for ax, img, lbl in zip(axes, imgs, raw_labels):
+    ax.imshow(img, cmap="magma", origin="lower", vmin=0, vmax=1)
+    ax.set_title(f"$\\Omega_m={lbl[0]:.2f}$, $\\sigma_8={lbl[1]:.2f}$", fontsize=10)
+    ax.set_xticks([]); ax.set_yticks([])
+plt.suptitle("Conditional DDPM samples — CAMELS HI fields", fontweight="bold")
+plt.tight_layout()
+plt.savefig("inference_example.png", dpi=150, bbox_inches="tight")
+print(f"\nSaved -> inference_example.png")

sample_modeling_ddpm_camels.py

@@ -0,0 +1,107 @@
+"""
+modeling_ddpm_camels.py
+=======================
+Self-contained loader for the conditional DDPM checkpoint hosted on the Hub.
+Users only need this file + diffusion_conditional.py + unet_conditional.py
++ config.json + model.safetensors  to run inference.
+"""
+from __future__ import annotations
+import json
+from pathlib import Path
+from typing import Dict, Tuple, Union
+
+import torch
+
+from diffusion_conditional import GaussianDiffusion, ConditionalDiffusionModel
+from unet_conditional import ConditionalUNet
+
+
+def build_model(config: Dict) -> ConditionalDiffusionModel:
+    """Instantiate the architecture from a config dict."""
+    unet = ConditionalUNet(
+        in_channels=int(config["in_channels"]),
+        out_channels=int(config["out_channels"]),
+        label_dim=int(config["label_dim"]),
+        base_channels=int(config["base_channels"]),
+        channel_multipliers=list(config["channel_multipliers"]),
+        attention_levels=list(config["attention_levels"]),
+        dropout=float(config["dropout"]),
+    )
+    diffusion = GaussianDiffusion(
+        timesteps=int(config["timesteps"]),
+        beta_start=float(config["beta_start"]),
+        beta_end=float(config["beta_end"]),
+        schedule_type=str(config["schedule_type"]),
+    )
+    return ConditionalDiffusionModel(unet, diffusion)
+
+
+def load_pretrained(
+    model_dir: Union[str, Path],
+    device: str = "cuda",
+) -> Tuple[ConditionalDiffusionModel, Dict]:
+    """
+    Load the model and its config from a directory containing:
+      - config.json
+      - model.safetensors  (or pytorch_model.bin as fallback)
+    """
+    model_dir = Path(model_dir)
+    config = json.loads((model_dir / "config.json").read_text())
+
+    model = build_model(config).to(device)
+
+    safetensors_path = model_dir / "model.safetensors"
+    bin_path         = model_dir / "pytorch_model.bin"
+    if safetensors_path.exists():
+        from safetensors.torch import load_file
+        state_dict = load_file(str(safetensors_path), device=device)
+    elif bin_path.exists():
+        state_dict = torch.load(bin_path, map_location=device, weights_only=True)
+    else:
+        raise FileNotFoundError(f"No model weights in {model_dir}")
+
+    # Allow partial-match loading for backward compatibility
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if missing:
+        print(f"  Warning: missing keys: {missing[:5]}{'...' if len(missing) > 5 else ''}")
+    if unexpected:
+        print(f"  Warning: unexpected keys: {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}")
+
+    model.eval()
+    for p in model.parameters():
+        p.requires_grad_(False)
+
+    return model, config
+
+
+# Convenience for one-shot inference
+def generate(
+    model: ConditionalDiffusionModel,
+    config: Dict,
+    raw_labels: torch.Tensor,        # (B, label_dim) — un-normalised cosmological params
+    n_samples: int = 1,
+    use_ddim: bool = True,
+    ddim_steps: int = None,
+    device: str = "cuda",
+) -> torch.Tensor:
+    """
+    Generate samples conditioned on raw (un-normalised) parameter values.
+
+    Returns: tensor of shape (B*n_samples, 1, H, W) in [-1, 1] model space.
+    """
+    if ddim_steps is None:
+        ddim_steps = config["ddim_steps_default"]
+
+    label_mu  = torch.tensor(config["label_mu"],  dtype=torch.float32, device=device)
+    label_std = torch.tensor(config["label_std"], dtype=torch.float32, device=device)
+
+    raw_labels = raw_labels.to(device)
+    norm_labels = (raw_labels - label_mu) / label_std
+    norm_labels = norm_labels.repeat_interleave(n_samples, dim=0)
+
+    H = W = config["image_size"]
+    return model.sample(
+        labels=norm_labels, channels=1, height=H, width=W,
+        use_ddim=use_ddim, ddim_steps=ddim_steps,
+        progress=False, device=device,
+    )