update readme

README.md

@@ -142,8 +142,8 @@ Notes:
 ### Basic Usage
 
 ```python
-import os
 from pathlib import Path
+import importlib.util
 import torch
 import torchaudio
 from transformers import AutoModel, AutoProcessor
@@ -159,6 +159,28 @@ pretrained_model_name_or_path = "OpenMOSS-Team/MOSS-SoundEffect"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16 if device == "cuda" else torch.float32
 
+def resolve_attn_implementation() -> str:
+    # Prefer FlashAttention 2 when package + device conditions are met.
+    if (
+        device == "cuda"
+        and importlib.util.find_spec("flash_attn") is not None
+        and dtype in {torch.float16, torch.bfloat16}
+    ):
+        major, _ = torch.cuda.get_device_capability()
+        if major >= 8:
+            return "flash_attention_2"
+
+    # CUDA fallback: use PyTorch SDPA kernels.
+    if device == "cuda":
+        return "sdpa"
+
+    # CPU fallback.
+    return "eager"
+
+
+attn_implementation = resolve_attn_implementation()
+print(f"[INFO] Using attn_implementation={attn_implementation}")
+
 processor = AutoProcessor.from_pretrained(
     pretrained_model_name_or_path,
     trust_remote_code=True,
@@ -176,14 +198,14 @@ conversations = [
 model = AutoModel.from_pretrained(
     pretrained_model_name_or_path,
     trust_remote_code=True,
-    attn_implementation="sdpa",
+    # If FlashAttention 2 is installed, you can set attn_implementation="flash_attention_2"
+    attn_implementation=attn_implementation,
     torch_dtype=dtype,
 ).to(device)
 model.eval()
 
 batch_size = 1
 
-messages = []
 save_dir = Path("inference_root")
 save_dir.mkdir(exist_ok=True, parents=True)
 sample_idx = 0
@@ -205,6 +227,7 @@ with torch.no_grad():
             out_path = save_dir / f"sample{sample_idx}.wav"
             sample_idx += 1
             torchaudio.save(out_path, audio.unsqueeze(0), processor.model_config.sampling_rate)
+
 ```
 
 ### Input Types