Bootstrap Audio Flamingo 3 custom endpoint repo

README.md

@@ -36,21 +36,23 @@ Then deploy a Dedicated Endpoint from that model repo.
 Important: make sure your endpoint repo contains top-level:
 - `handler.py`
 - `requirements.txt`
-- `sitecustomize.py`
-- `setup.py`
+- `README.md`
 
-from this folder.
-If logs say `No custom pipeline found at /repository/handler.py`, your files were not copied to repo root.
-Use Endpoint task `custom` so the runtime loads `handler.py` instead of the default Transformers pipeline.
+Use endpoint task `custom` so the runtime loads `handler.py` instead of a default Transformers pipeline.
 
-## Endpoint env var
+## Endpoint env vars
 
+Required:
 - `AF3_MODEL_ID=nvidia/audio-flamingo-3-hf`
-- `PYTHONPATH=/repository` (ensures `sitecustomize.py` compatibility patch is loaded)
+
+Optional runtime bootstrap (defaults shown):
+- `AF3_BOOTSTRAP_RUNTIME=1`
+- `AF3_TRANSFORMERS_SPEC=transformers==5.1.0`
+- `AF3_RUNTIME_DIR=/tmp/af3_runtime`
+- `AF3_STUB_TORCHVISION=1`
 
 ## Notes
 
 - Audio Flamingo 3 is large; use a GPU endpoint.
-- This handler returns raw prose analysis. Use the local AF3+ChatGPT pipeline to normalize to LoRA sidecar JSON.
-- If logs show `cannot import name 'is_tf_available' from transformers.file_utils`,
-  ensure `sitecustomize.py` is present in repo root and endpoint env includes `PYTHONPATH=/repository`.
+- First boot can take longer because the handler installs AF3-compatible runtime dependencies.
+- This handler returns raw prose analysis. Use the local AF3+ChatGPT pipeline to normalize to LoRA sidecar JSON.

handler.py

@@ -1,11 +1,15 @@
 import base64
+import importlib
+import io
 import os
-import tempfile
-from typing import Any, Dict, List
+import subprocess
+import sys
+import types
+from typing import Any, Dict, List, Tuple
 
+import numpy as np
+import soundfile as sf
 import torch
-import transformers
-from transformers import AutoProcessor
 
 
 def _resolve_model_id(model_dir: str) -> str:
@@ -17,6 +21,134 @@ def _resolve_model_id(model_dir: str) -> str:
     return default_id
 
 
+def _log(msg: str) -> None:
+    print(f"[AF3 handler] {msg}", flush=True)
+
+
+def _env_true(name: str, default: bool = False) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    return str(raw).strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _install_torchvision_stub() -> None:
+    if not _env_true("AF3_STUB_TORCHVISION", True):
+        return
+    if "torchvision" in sys.modules:
+        return
+    interpolation_mode = types.SimpleNamespace(
+        NEAREST=0,
+        BILINEAR=2,
+        BICUBIC=3,
+        BOX=4,
+        HAMMING=5,
+        LANCZOS=1,
+    )
+    transforms_stub = types.ModuleType("torchvision.transforms")
+    setattr(transforms_stub, "InterpolationMode", interpolation_mode)
+    tv_stub = types.ModuleType("torchvision")
+    setattr(tv_stub, "transforms", transforms_stub)
+    sys.modules["torchvision"] = tv_stub
+    sys.modules["torchvision.transforms"] = transforms_stub
+
+
+def _clear_python_modules(prefixes: Tuple[str, ...]) -> None:
+    for name in list(sys.modules.keys()):
+        if any(name == p or name.startswith(f"{p}.") for p in prefixes):
+            sys.modules.pop(name, None)
+
+
+def _af3_classes_available() -> tuple[bool, str]:
+    try:
+        from transformers import AudioFlamingo3ForConditionalGeneration  # noqa: F401
+        from transformers import AudioFlamingo3Processor  # noqa: F401
+
+        return True, ""
+    except Exception as exc:
+        return False, f"{type(exc).__name__}: {exc}"
+
+
+def _bootstrap_runtime_transformers(target_dir: str) -> None:
+    packages = [
+        os.getenv("AF3_TRANSFORMERS_SPEC", "transformers==5.1.0"),
+        "numpy<2",
+        "accelerate>=1.1.0",
+        "sentencepiece",
+        "safetensors",
+        "soxr",
+    ]
+    cmd = [sys.executable, "-m", "pip", "install", "--upgrade", "--no-cache-dir", "--target", target_dir, *packages]
+    _log("Installing runtime deps for AF3 (first boot can take a few minutes).")
+    subprocess.check_call(cmd)
+
+
+def _ensure_af3_transformers():
+    _install_torchvision_stub()
+
+    import transformers
+
+    ok, err = _af3_classes_available()
+    if ok:
+        _log(f"Using bundled transformers={transformers.__version__}")
+        return transformers
+
+    if not _env_true("AF3_BOOTSTRAP_RUNTIME", True):
+        raise RuntimeError(
+            "AF3 classes are unavailable in bundled transformers "
+            f"({transformers.__version__}) and AF3_BOOTSTRAP_RUNTIME is disabled. "
+            f"Last import error: {err}"
+        )
+
+    target_dir = os.getenv("AF3_RUNTIME_DIR", "/tmp/af3_runtime")
+    os.makedirs(target_dir, exist_ok=True)
+    _bootstrap_runtime_transformers(target_dir)
+    if target_dir not in sys.path:
+        sys.path.insert(0, target_dir)
+
+    _clear_python_modules(("transformers", "tokenizers", "huggingface_hub", "safetensors"))
+    _install_torchvision_stub()
+    importlib.invalidate_caches()
+    transformers = importlib.import_module("transformers")
+
+    ok, err = _af3_classes_available()
+    if not ok:
+        raise RuntimeError(
+            "Failed to load AF3 processor classes after runtime bootstrap. "
+            f"transformers={getattr(transformers, '__version__', 'unknown')} "
+            f"error={err}"
+        )
+    _log(f"Bootstrapped transformers={transformers.__version__}")
+    return transformers
+
+
+def _resample_audio_mono(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
+    if src_sr == dst_sr:
+        return audio.astype(np.float32, copy=False)
+    if audio.size == 0:
+        return np.zeros((0,), dtype=np.float32)
+    src_idx = np.arange(audio.shape[0], dtype=np.float64)
+    dst_len = int(round(audio.shape[0] * float(dst_sr) / float(src_sr)))
+    dst_len = max(dst_len, 1)
+    dst_idx = np.linspace(0.0, float(max(audio.shape[0] - 1, 0)), dst_len, dtype=np.float64)
+    out = np.interp(dst_idx, src_idx, audio.astype(np.float64, copy=False))
+    return out.astype(np.float32, copy=False)
+
+
+def _decode_audio_from_b64(audio_b64: str) -> tuple[np.ndarray, int]:
+    raw = base64.b64decode(audio_b64)
+    data, sr = sf.read(io.BytesIO(raw), dtype="float32", always_2d=False)
+    if data.ndim == 2:
+        data = np.mean(data, axis=1)
+    if data.ndim != 1:
+        data = np.asarray(data).reshape(-1)
+    target_sr = 16000
+    if int(sr) != target_sr:
+        data = _resample_audio_mono(data, int(sr), target_sr)
+        sr = target_sr
+    return data.astype(np.float32, copy=False), int(sr)
+
+
 class EndpointHandler:
     """
     Hugging Face Dedicated Endpoint custom handler.
@@ -37,30 +169,17 @@ class EndpointHandler:
 
     def __init__(self, model_dir: str = ""):
         self.model_id = _resolve_model_id(model_dir)
-        print(
-            f"[AF3 handler] transformers={transformers.__version__} "
-            f"AudioFlamingo3Processor={hasattr(transformers, 'AudioFlamingo3Processor')} "
-            f"AudioFlamingo3ForConditionalGeneration={hasattr(transformers, 'AudioFlamingo3ForConditionalGeneration')}",
-            flush=True,
+        self.transformers = _ensure_af3_transformers()
+        from transformers import AudioFlamingo3ForConditionalGeneration, AutoProcessor
+
+        _log(
+            f"torch={torch.__version__} cuda={torch.cuda.is_available()} "
+            f"transformers={self.transformers.__version__} model_id={self.model_id}"
         )
-        try:
-            from transformers import AudioFlamingo3ForConditionalGeneration
-            model_cls = AudioFlamingo3ForConditionalGeneration
-        except Exception:
-            from transformers import AutoModelForImageTextToText
-            model_cls = AutoModelForImageTextToText
 
         dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-        try:
-            self.processor = AutoProcessor.from_pretrained(self.model_id, trust_remote_code=True)
-        except Exception as exc:
-            raise RuntimeError(
-                "Failed to load AF3 processor. "
-                f"transformers={transformers.__version__} "
-                f"AudioFlamingo3Processor={hasattr(transformers, 'AudioFlamingo3Processor')} "
-                f"model_id={self.model_id} error={exc}"
-            ) from exc
-        self.model = model_cls.from_pretrained(
+        self.processor = AutoProcessor.from_pretrained(self.model_id, trust_remote_code=True)
+        self.model = AudioFlamingo3ForConditionalGeneration.from_pretrained(
             self.model_id,
             torch_dtype=dtype,
             trust_remote_code=True,
@@ -68,41 +187,31 @@ class EndpointHandler:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model.to(self.device)
 
-    def _build_inputs(self, audio_path: str, prompt: str) -> Dict[str, Any]:
-        conversation_variants: List[List[Dict[str, Any]]] = [
-            [
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "audio", "path": audio_path},
-                        {"type": "text", "text": prompt},
-                    ],
-                }
-            ],
-            [
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "audio", "audio_url": audio_path},
-                        {"type": "text", "text": prompt},
-                    ],
-                }
-            ],
+    def _build_inputs(self, audio: np.ndarray, sample_rate: int, prompt: str) -> Dict[str, Any]:
+        conversation: List[Dict[str, Any]] = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "audio", "audio": audio},
+                    {"type": "text", "text": prompt},
+                ],
+            }
         ]
-
-        last_exc: Exception | None = None
-        for convo in conversation_variants:
-            try:
-                return self.processor.apply_chat_template(
-                    convo,
-                    tokenize=True,
-                    add_generation_prompt=True,
-                    return_dict=True,
-                )
-            except Exception as exc:
-                last_exc = exc
-                continue
-        raise RuntimeError(f"Failed to build AF3 inputs from chat template: {last_exc}")
+        try:
+            return self.processor.apply_chat_template(
+                conversation,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_dict=True,
+                audio_kwargs={"sampling_rate": int(sample_rate)},
+            )
+        except Exception:
+            return self.processor.apply_chat_template(
+                conversation,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_dict=True,
+            )
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         payload = data.get("inputs", data) if isinstance(data, dict) else {}
@@ -114,14 +223,9 @@ class EndpointHandler:
         max_new_tokens = int(payload.get("max_new_tokens", 1200))
         temperature = float(payload.get("temperature", 0.1))
 
-        tmp_path = ""
         try:
-            raw = base64.b64decode(audio_b64)
-            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-                tmp.write(raw)
-                tmp_path = tmp.name
-
-            inputs = self._build_inputs(tmp_path, prompt)
+            audio, sample_rate = _decode_audio_from_b64(audio_b64)
+            inputs = self._build_inputs(audio, sample_rate, prompt)
             device = next(self.model.parameters()).device
             for key, value in list(inputs.items()):
                 if hasattr(value, "to"):
@@ -144,10 +248,4 @@ class EndpointHandler:
                 text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()
             return {"generated_text": text}
         except Exception as exc:
-            return {"error": str(exc)}
-        finally:
-            if tmp_path:
-                try:
-                    os.unlink(tmp_path)
-                except Exception:
-                    pass
+            return {"error": str(exc)}

requirements.txt

@@ -1,12 +1,2 @@
--e .
-torch
-torchaudio
-soundfile
-numpy
-transformers==5.0.0rc1
-huggingface_hub>=1.0.0
-accelerate>=1.0.0
-diffusers>=0.35.0
-peft>=0.17.0
-sentencepiece
-safetensors
+numpy<2
+soundfile