Deploy BYOD-Llama-3.1-8B full-precision demo

app.py

@@ -25,6 +25,8 @@ SESSION = load_hub_adapter_session(
     MODEL_REPO_ID,
     device_name="cuda",
     quantization="none",
+    # A forward pass is only valid after @spaces.GPU has allocated hardware.
+    preflight=False,
 )
 print(f"Loaded {DISPLAY_NAME} ({SESSION.compute_dtype}, unquantized).")
 

src/diffusion_lm/inference.py

@@ -82,6 +82,7 @@ def _load_adapter_path(
     adapter_path: str | Path,
     device_name: str = "auto",
     quantization: str | None = None,
+    preflight: bool = True,
 ) -> InferenceSession:
     """Load an adapter directory that has already been resolved and validated."""
     adapter_path = Path(adapter_path).expanduser().resolve()
@@ -166,7 +167,8 @@ def _load_adapter_path(
     model.eval()
     mask_info = validate_mask_token(tokenizer, str(run_config.get("mask_token", "MASK")))
     session = InferenceSession(model, tokenizer, device, adapter_path, run_config, mask_info["mask_token_id"], "4bit" if use_4bit else "none", str(compute_dtype).removeprefix("torch."))
-    preflight_session(session)
+    if preflight:
+        preflight_session(session)
     return session
 
 
@@ -182,6 +184,7 @@ def load_hub_adapter_session(
     quantization: str | None = None,
     revision: str | None = None,
     cache_dir: str | Path | None = None,
+    preflight: bool = True,
 ) -> InferenceSession:
     """Download and load a BYOD adapter from the Hugging Face Hub."""
     try:
@@ -197,7 +200,7 @@ def load_hub_adapter_session(
         cache_dir=str(cache_dir) if cache_dir is not None else None,
         token=os.getenv("HF_TOKEN"),
     )
-    return _load_adapter_path(adapter_path, device_name, quantization)
+    return _load_adapter_path(adapter_path, device_name, quantization, preflight=preflight)
 
 
 def load_merged_session(