Support Fish Speech TTS inference on XPU platform

docs/models/supported_models.md

@@ -82,7 +82,7 @@ th {
 | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
 | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
-| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
+| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ |
 | `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
 | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
 | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -122,12 +122,48 @@ def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
     logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sr)
 
 
+def _collect_audio_chunk(
+    mm: dict,
+    all_audio_chunks: list[torch.Tensor],
+    sample_rate: int | None,
+) -> tuple[int | None, int]:
+    """Extract audio and sample rate from a multimodal output chunk.
+
+    Returns updated sample_rate and the number of audio segments in this chunk.
+    """
+    audio = mm.get("audio") if mm else None
+    if audio is not None:
+        if isinstance(audio, list):
+            all_audio_chunks.extend(audio)
+        else:
+            all_audio_chunks.append(audio)
+    sr_raw = mm.get("sr") if mm else None
+    if sr_raw is not None and not (isinstance(sr_raw, (list, tuple)) and len(sr_raw) == 0):
+        sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
+        sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+    n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
+    return sample_rate, n
+
+
+def _build_omni_kwargs(args, model_name: str) -> dict:
+    """Build keyword arguments for Omni/AsyncOmni initialization."""
+    omni_kwargs = dict(
+        model=model_name,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
+    else:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
+    return omni_kwargs
+
+
 def main(args):
     """Run offline inference with Omni."""
     output_dir = args.output_dir
     os.makedirs(output_dir, exist_ok=True)
 
-    stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG
     model_name = args.model or DEFAULT_MODEL
 
     inputs = [
@@ -139,12 +175,7 @@ def main(args):
         )
     ]
 
-    omni = Omni(
-        model=model_name,
-        stage_configs_path=stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-    )
+    omni = Omni(**_build_omni_kwargs(args, model_name))
 
     t_start = time.perf_counter()
     for stage_outputs in omni.generate(inputs):
@@ -165,7 +196,6 @@ async def main_streaming(args):
     output_dir = args.output_dir
     os.makedirs(output_dir, exist_ok=True)
 
-    stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG
     model_name = args.model or DEFAULT_MODEL
 
     prompt = build_prompt(
@@ -175,24 +205,20 @@ async def main_streaming(args):
         model_name=model_name,
     )
 
-    omni = AsyncOmni(
-        model=model_name,
-        stage_configs_path=stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-    )
+    omni = AsyncOmni(**_build_omni_kwargs(args, model_name))
 
     request_id = "0"
     t_start = time.perf_counter()
     t_prev = t_start
     chunk_idx = 0
+    all_audio_chunks: list[torch.Tensor] = []
+    sample_rate = None
 
     async for stage_output in omni.generate(prompt, request_id=request_id):
         mm = stage_output.request_output.outputs[0].multimodal_output
         if not stage_output.finished:
             t_now = time.perf_counter()
-            audio = mm.get("audio")
-            n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
+            sample_rate, n = _collect_audio_chunk(mm, all_audio_chunks, sample_rate)
             dt_ms = (t_now - t_prev) * 1000
             ttfa_ms = (t_now - t_start) * 1000
             if chunk_idx == 0:
@@ -202,10 +228,20 @@ async def main_streaming(args):
             t_prev = t_now
             chunk_idx += 1
         else:
+            # Collect any remaining audio in the final message
+            sample_rate, _ = _collect_audio_chunk(mm, all_audio_chunks, sample_rate)
+
             t_end = time.perf_counter()
             total_ms = (t_end - t_start) * 1000
             logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx)
-            _save_wav(output_dir, request_id, mm)
+
+            if all_audio_chunks and sample_rate:
+                audio_tensor = torch.cat(all_audio_chunks, dim=-1)
+                out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+                sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sample_rate, format="WAV")
+                logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sample_rate)
+            else:
+                logger.warning("Request %s: no audio chunks received", request_id)
 
 
 def parse_args():
@@ -236,11 +272,17 @@ def parse_args():
         default=None,
         help=f"HuggingFace model path (default: {DEFAULT_MODEL}).",
     )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help="Deploy config YAML path (under vllm_omni/deploy/).",
+    )
     parser.add_argument(
         "--stage-configs-path",
         type=str,
         default=None,
-        help="Path to stage configs YAML.",
+        help="Path to stage configs YAML (legacy, prefer --deploy-config).",
     )
     parser.add_argument(
         "--output-dir",

vllm_omni/attention/fish_kvcache_attn.py

@@ -111,6 +111,8 @@ def can_use_fish_kvcache_attn(
 
 
 def _is_cuda_graph_capturing() -> bool:
+    if not torch.cuda.is_available():
+        return False
     return bool(torch.cuda.is_current_stream_capturing())
 
 

vllm_omni/deploy/fish_qwen3_omni.yaml

@@ -63,3 +63,15 @@ stages:
       max_tokens: 65536
       seed: 42
       repetition_penalty: 1.0
+
+platforms:
+  xpu:
+    stages:
+      # Verified on Intel Arc B60.
+      - stage_id: 0
+        gpu_memory_utilization: 0.8
+        max_model_len: 2048
+        max_num_batched_tokens: 2048
+      - stage_id: 1
+        max_model_len: 2048
+        max_num_batched_tokens: 2048