Bootstrap Audio Flamingo 3 NVIDIA-stack endpoint repo

README.md

@@ -0,0 +1,54 @@
+# Audio Flamingo 3 NVIDIA-Stack Endpoint Template
+
+This template uses the same core runtime pattern as NVIDIA's Space:
+- `llava` code from `nvidia/audio-flamingo-3` (space repo)
+- base checkpoint from `nvidia/audio-flamingo-3` (model repo)
+- optional `stage35` think/long adapter
+
+## Request contract
+
+```json
+{
+  "inputs": {
+    "prompt": "Please describe the audio in detail.",
+    "audio_base64": "<base64 WAV bytes>",
+    "think_mode": true,
+    "max_new_tokens": 2048,
+    "temperature": 0.2
+  }
+}
+```
+
+## Response contract
+
+```json
+{
+  "generated_text": "...",
+  "mode": "think"
+}
+```
+
+## Bootstrap command
+
+```bash
+python scripts/hf_clone.py af3-nvidia-endpoint --repo-id YOUR_USERNAME/YOUR_AF3_NVIDIA_ENDPOINT_REPO
+```
+
+## Endpoint settings
+
+- Task: `custom`
+- GPU instance required
+- Secrets:
+  - `HF_TOKEN=<your_token>`
+
+## Optional env vars
+
+- `AF3_NV_CODE_REPO_ID=nvidia/audio-flamingo-3`
+- `AF3_NV_MODEL_REPO_ID=nvidia/audio-flamingo-3`
+- `AF3_NV_CODE_REPO_TYPE=space`
+- `AF3_NV_MODEL_REPO_TYPE=model`
+- `AF3_NV_DEFAULT_MODE=think`
+- `AF3_NV_LOAD_THINK=1`
+- `AF3_NV_LOAD_SINGLE=0`
+
+Default behavior loads think/long mode for higher-quality long-form reasoning.

handler.py

@@ -0,0 +1,204 @@
+import base64
+import copy
+import os
+import sys
+import tempfile
+from typing import Any, Dict
+
+import torch
+from huggingface_hub import snapshot_download
+from peft import PeftModel
+
+
+DEFAULT_PROMPT = "Please describe the audio in detail."
+
+
+def _log(msg: str) -> None:
+    print(f"[AF3 NVIDIA 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 _strip_state_dict_prefixes(state_dict: Dict[str, Any]) -> Dict[str, Any]:
+    out: Dict[str, Any] = {}
+    for key, value in state_dict.items():
+        key2 = key[6:] if key.startswith("model.") else key
+        out[key2] = value
+    return out
+
+
+class EndpointHandler:
+    """
+    NVIDIA AF3 stack endpoint handler (matches Space architecture closely).
+
+    Request:
+      {
+        "inputs": {
+          "prompt": "...",
+          "audio_base64": "...",
+          "think_mode": true,
+          "max_new_tokens": 2048,
+          "temperature": 0.2
+        }
+      }
+
+    Response:
+      {"generated_text": "...", "mode": "think|single"}
+    """
+
+    def __init__(self, model_dir: str = ""):
+        del model_dir
+        self.hf_token = os.getenv("HF_TOKEN", "")
+        self.code_repo_id = os.getenv("AF3_NV_CODE_REPO_ID", "nvidia/audio-flamingo-3")
+        self.model_repo_id = os.getenv("AF3_NV_MODEL_REPO_ID", "nvidia/audio-flamingo-3")
+        self.code_repo_type = os.getenv("AF3_NV_CODE_REPO_TYPE", "space")
+        self.model_repo_type = os.getenv("AF3_NV_MODEL_REPO_TYPE", "model")
+        self.default_mode = os.getenv("AF3_NV_DEFAULT_MODE", "think").strip().lower()
+        if self.default_mode not in {"think", "single"}:
+            self.default_mode = "think"
+
+        self.load_think = _env_true("AF3_NV_LOAD_THINK", True)
+        self.load_single = _env_true("AF3_NV_LOAD_SINGLE", self.default_mode == "single")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        _log(f"torch={torch.__version__} cuda={torch.cuda.is_available()} device={self.device}")
+        _log(
+            f"code_repo={self.code_repo_type}:{self.code_repo_id} "
+            f"model_repo={self.model_repo_type}:{self.model_repo_id} default_mode={self.default_mode}"
+        )
+
+        self.llava = self._load_llava_runtime()
+        self.model_root = self._download_model_root()
+
+        self.model_single = None
+        self.model_think = None
+
+        if self.load_single:
+            self.model_single = self._load_single_model()
+        if self.load_think:
+            self.model_think = self._load_think_model()
+
+        if self.model_single is None and self.model_think is None:
+            raise RuntimeError("No model loaded. Enable AF3_NV_LOAD_THINK or AF3_NV_LOAD_SINGLE.")
+
+    def _load_llava_runtime(self):
+        code_root = snapshot_download(
+            repo_id=self.code_repo_id,
+            repo_type=self.code_repo_type,
+            allow_patterns=["llava/**"],
+            token=self.hf_token or None,
+        )
+        if code_root not in sys.path:
+            sys.path.insert(0, code_root)
+        import llava  # type: ignore
+
+        _log(f"Loaded llava runtime from {code_root}")
+        return llava
+
+    def _download_model_root(self) -> str:
+        model_root = snapshot_download(
+            repo_id=self.model_repo_id,
+            repo_type=self.model_repo_type,
+            token=self.hf_token or None,
+        )
+        _log(f"Model root: {model_root}")
+        return model_root
+
+    def _load_single_model(self):
+        _log("Loading single-turn model...")
+        model = self.llava.load(self.model_root, model_base=None)
+        model = model.to(self.device)
+        model.eval()
+        return model
+
+    def _load_think_model(self):
+        _log("Loading think/long model (stage35 adapter)...")
+        stage35_dir = os.path.join(self.model_root, "stage35")
+        non_lora_path = os.path.join(stage35_dir, "non_lora_trainables.bin")
+        if not os.path.exists(non_lora_path):
+            raise RuntimeError(f"stage35 non_lora_trainables missing: {non_lora_path}")
+
+        model = self.llava.load(self.model_root, model_base=None)
+        model = model.to(self.device)
+
+        non_lora_trainables = torch.load(non_lora_path, map_location="cpu")
+        non_lora_trainables = _strip_state_dict_prefixes(non_lora_trainables)
+        model.load_state_dict(non_lora_trainables, strict=False)
+
+        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        model = PeftModel.from_pretrained(
+            model,
+            stage35_dir,
+            device_map="auto" if torch.cuda.is_available() else None,
+            torch_dtype=dtype,
+        )
+        model.eval()
+        return model
+
+    def _select_model(self, think_mode: bool):
+        if think_mode and self.model_think is not None:
+            return self.model_think, "think"
+        if (not think_mode) and self.model_single is not None:
+            return self.model_single, "single"
+        if self.model_think is not None:
+            return self.model_think, "think"
+        return self.model_single, "single"
+
+    def _build_generation_config(self, model, max_new_tokens: int, temperature: float):
+        base_cfg = getattr(model, "default_generation_config", None)
+        if base_cfg is None:
+            return None
+        cfg = copy.deepcopy(base_cfg)
+        if max_new_tokens > 0:
+            setattr(cfg, "max_new_tokens", int(max_new_tokens))
+        if temperature > 0:
+            setattr(cfg, "temperature", float(temperature))
+            setattr(cfg, "do_sample", True)
+        else:
+            setattr(cfg, "do_sample", False)
+        return cfg
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        payload = data.get("inputs", data) if isinstance(data, dict) else {}
+        audio_b64 = payload.get("audio_base64")
+        if not audio_b64:
+            return {"error": "audio_base64 is required"}
+
+        prompt = str(payload.get("prompt", DEFAULT_PROMPT)).strip() or DEFAULT_PROMPT
+        think_mode_val = payload.get("think_mode")
+        if think_mode_val is None:
+            think_mode = self.default_mode == "think"
+        else:
+            think_mode = bool(think_mode_val)
+
+        max_new_tokens = int(payload.get("max_new_tokens", 2048))
+        temperature = float(payload.get("temperature", 0.2))
+        model, mode = self._select_model(think_mode)
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+            tmp_path = tmp.name
+            tmp.write(base64.b64decode(audio_b64))
+
+        try:
+            sound = self.llava.Sound(tmp_path)
+            full_prompt = f"<sound>\n{prompt}"
+            gen_cfg = self._build_generation_config(model, max_new_tokens=max_new_tokens, temperature=temperature)
+
+            with torch.inference_mode():
+                if gen_cfg is not None:
+                    response = model.generate_content([sound, full_prompt], generation_config=gen_cfg)
+                else:
+                    response = model.generate_content([sound, full_prompt])
+            return {"generated_text": str(response).strip(), "mode": mode}
+        except Exception as exc:
+            return {"error": str(exc), "mode": mode}
+        finally:
+            try:
+                os.unlink(tmp_path)
+            except Exception:
+                pass

requirements.txt

@@ -0,0 +1,19 @@
+transformers==4.46.0
+accelerate==0.34.2
+peft==0.14.0
+numpy==1.26.4
+Pillow
+pydub
+soundfile
+librosa
+openai-whisper
+ftfy
+jiwer
+einops
+hydra-core
+loguru
+pytorchvideo==0.1.5
+opencv-python-headless==4.8.0.76
+protobuf==3.20.*
+termcolor
+sentencepiece