Upload inference.py

inference.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""Simple Hugging Face inference script for the digit diffusion model.
+
+No command-line arguments. Edit the values in the CONFIG section below.
+
+What it does:
+- loads the model from the Hugging Face Hub (or a local HF cache/path)
+- loads the DDPM scheduler from the same repo
+- generates one or more images for one digit or several digits
+- saves everything as a single PNG grid
+"""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Iterable
+
+import torch
+from diffusers import DDPMScheduler
+from torchvision.utils import make_grid, save_image
+from transformers import AutoModel
+
+
+# -----------------------------------------------------------------------------
+# CONFIG — edit these values only
+# -----------------------------------------------------------------------------
+
+MODEL_ID = "your-hf-username/your-digit-diffusion-repo"
+OUTPUT_IMAGE = "./digit_samples.png"
+
+# Choose either a single digit or multiple digits.
+USE_MULTIPLE_DIGITS = False
+DIGIT = 7
+DIGITS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+
+# How many images to generate for each selected digit.
+IMAGES_PER_DIGIT = 4
+
+# Number of denoising steps.
+NUM_INFERENCE_STEPS = 1000
+
+# Output image size should match training.
+IMAGE_SIZE = 32
+
+# Reproducibility.
+SEED = 42
+
+# Optional performance knobs.
+USE_AMP = torch.cuda.is_available()
+ALLOW_TF32 = True
+
+
+# -----------------------------------------------------------------------------
+# Helpers
+# -----------------------------------------------------------------------------
+
+def _selected_digits() -> list[int]:
+    if USE_MULTIPLE_DIGITS:
+        if not DIGITS:
+            raise ValueError("DIGITS must not be empty when USE_MULTIPLE_DIGITS=True")
+        return [int(d) for d in DIGITS]
+    return [int(DIGIT)]
+
+
+def _load_model(model_id: str, device: torch.device):
+    """Load the custom HF model without defining any local model classes."""
+    model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
+    model.to(device)
+    model.eval()
+    return model
+
+
+def _load_scheduler(model_id: str) -> DDPMScheduler:
+    return DDPMScheduler.from_pretrained(model_id)
+
+
+def _to_display_range(x: torch.Tensor) -> torch.Tensor:
+    """Map tensors from [-1, 1] to [0, 1]."""
+    return ((x.clamp(-1.0, 1.0) + 1.0) / 2.0).cpu()
+
+
+@torch.inference_mode()
+def generate_grid(
+    model,
+    scheduler: DDPMScheduler,
+    device: torch.device,
+    digits: Iterable[int],
+    images_per_digit: int,
+    num_inference_steps: int,
+    image_size: int,
+) -> torch.Tensor:
+    """Generate a single grid image containing all requested digits."""
+    scheduler.set_timesteps(num_inference_steps, device=device)
+
+    rows: list[torch.Tensor] = []
+    for digit in digits:
+        autocast_ctx = (
+            torch.autocast(device_type="cuda", dtype=torch.float16)
+            if USE_AMP and device.type == "cuda"
+            else nullcontext()
+        )
+        latents = torch.randn(
+            images_per_digit,
+            1,
+            image_size,
+            image_size,
+            device=device,
+        )
+        class_labels = torch.full(
+            (images_per_digit,),
+            int(digit),
+            device=device,
+            dtype=torch.long,
+        )
+
+        with autocast_ctx:
+            for t in scheduler.timesteps:
+                t_batch = torch.full(
+                    (images_per_digit,),
+                    int(t),
+                    device=device,
+                    dtype=torch.long,
+                )
+
+                output = model(
+                    noisy_images=latents,
+                    timesteps=t_batch,
+                    class_labels=class_labels,
+                )
+                noise_pred = output.sample if hasattr(output, "sample") else output[0]
+                latents = scheduler.step(noise_pred, t, latents).prev_sample
+
+        rows.append(_to_display_range(latents))
+
+    all_images = torch.cat(rows, dim=0)
+    nrow = images_per_digit
+    grid = make_grid(all_images, nrow=nrow)
+    return grid
+
+
+# -----------------------------------------------------------------------------
+# Main
+# -----------------------------------------------------------------------------
+
+def main() -> None:
+    if ALLOW_TF32 and torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    torch.manual_seed(SEED)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(SEED)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    digits = _selected_digits()
+
+    print(f"[info] model_id          = {MODEL_ID}")
+    print(f"[info] device            = {device}")
+    print(f"[info] digits            = {digits}")
+    print(f"[info] images_per_digit  = {IMAGES_PER_DIGIT}")
+    print(f"[info] num_steps         = {NUM_INFERENCE_STEPS}")
+    print(f"[info] output_image      = {OUTPUT_IMAGE}")
+
+    model = _load_model(MODEL_ID, device)
+    scheduler = _load_scheduler(MODEL_ID)
+
+    grid = generate_grid(
+        model=model,
+        scheduler=scheduler,
+        device=device,
+        digits=digits,
+        images_per_digit=IMAGES_PER_DIGIT,
+        num_inference_steps=NUM_INFERENCE_STEPS,
+        image_size=IMAGE_SIZE,
+    )
+
+    out_path = Path(OUTPUT_IMAGE)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    save_image(grid, out_path)
+    print(f"[done] saved -> {out_path.resolve()}")
+
+
+if __name__ == "__main__":
+    main()