Update inference.py

inference.py

@@ -23,10 +23,6 @@ USE_AMP = torch.cuda.is_available()
 ALLOW_TF32 = False
 
 
-# -----------------------------------------------------------------------------
-# Helpers
-# -----------------------------------------------------------------------------
-
 def _selected_digits() -> list[int]:
     if USE_MULTIPLE_DIGITS:
         if not DIGITS:
@@ -43,7 +39,7 @@ def _load_model(model_id: str, device: torch.device):
 
 
 def _load_scheduler(model_id: str) -> DDPMScheduler:
-    return DDPMScheduler.from_pretrained(model_id)
+    return DDPMScheduler.from_pretrained(model_id, trust_remote_code=True)
 
 
 def _to_display_range(x: torch.Tensor) -> torch.Tensor:
@@ -60,7 +56,6 @@ def generate_grid(
     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] = []
@@ -108,6 +103,7 @@ def generate_grid(
     grid = make_grid(all_images, nrow=nrow)
     return grid
 
+
 def main() -> None:
     if ALLOW_TF32 and torch.cuda.is_available():
         torch.backends.cuda.matmul.allow_tf32 = True
@@ -146,4 +142,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main()