Update inference.py

inference.py

@@ -1,15 +1,4 @@
 #!/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
@@ -22,49 +11,31 @@ 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"
+# config
 
-# Choose either a single digit or multiple digits.
+MODEL_ID = "Harley-ml/MNIST-IMG-390k"
+OUTPUT_IMAGE = "./digits.png"
 USE_MULTIPLE_DIGITS = False
-DIGIT = 7
+DIGIT = 1
 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
+ALLOW_TF32 = False
 
-
-# -----------------------------------------------------------------------------
 # 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")
+            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()
@@ -76,7 +47,6 @@ def _load_scheduler(model_id: str) -> DDPMScheduler:
 
 
 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()
 
 
@@ -177,7 +147,7 @@ def main() -> None:
     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()}")
+    print(f"[done] saved to {out_path.resolve()}")
 
 
 if __name__ == "__main__":