""" Teste de Stress - Validar capacidade real de usuarios simultaneos """ import os os.environ["VLLM_ATTENTION_BACKEND"] = "FLASH_ATTN" import torch import time import asyncio import numpy as np from transformers import AutoTokenizer from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams from snac import SNAC START_TOKEN = 128259 END_TOKENS = [128009, 128260, 128261, 128257] STOP_TOKEN = 128258 AUDIO_TOKEN_BASE = 128266 def decode_tokens_to_audio(token_ids, snac_model): audio_frames = [] for t in token_ids: if isinstance(t, str): continue if t >= AUDIO_TOKEN_BASE: offset = t - AUDIO_TOKEN_BASE layer = offset // 4096 value = offset % 4096 audio_frames.append((layer, value)) if len(audio_frames) < 7: return None num_complete_frames = len(audio_frames) // 7 if num_complete_frames == 0: return None codes_0, codes_1, codes_2 = [], [], [] for i in range(num_complete_frames): base = i * 7 codes_0.append(audio_frames[base][1]) codes_1.append(audio_frames[base + 1][1]) codes_1.append(audio_frames[base + 4][1]) codes_2.append(audio_frames[base + 2][1]) codes_2.append(audio_frames[base + 3][1]) codes_2.append(audio_frames[base + 5][1]) codes_2.append(audio_frames[base + 6][1]) try: with torch.no_grad(): codes = [ torch.tensor(codes_0, dtype=torch.int64).unsqueeze(0).to("cuda"), torch.tensor(codes_1, dtype=torch.int64).unsqueeze(0).to("cuda"), torch.tensor(codes_2, dtype=torch.int64).unsqueeze(0).to("cuda"), ] audio = snac_model.decode(codes) return audio.squeeze().cpu().numpy() except: return None async def main(): print("=" * 70) print("TESTE DE STRESS - VALIDAR CAPACIDADE REAL DE USUARIOS") print("=" * 70) print("\n[1] Carregando modelo...") tokenizer = AutoTokenizer.from_pretrained("canopylabs/orpheus-3b-0.1-ft") engine_args = AsyncEngineArgs( model="canopylabs/orpheus-3b-0.1-ft", dtype="bfloat16", max_model_len=4096, gpu_memory_utilization=0.9, max_num_seqs=16, enable_chunked_prefill=True, enable_prefix_caching=True, enforce_eager=False, ) engine = AsyncLLMEngine.from_engine_args(engine_args) snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to("cuda") sampling_params = SamplingParams( temperature=0.2, top_p=0.9, max_tokens=4096, stop_token_ids=[STOP_TOKEN], repetition_penalty=1.1, ) def format_prompt(text, voice="tara"): adapted_prompt = f"{voice}: {text}" prompt_tokens = tokenizer(adapted_prompt, return_tensors="pt") start_token = torch.tensor([[START_TOKEN]], dtype=torch.int64) end_tokens = torch.tensor([END_TOKENS], dtype=torch.int64) all_input_ids = torch.cat([start_token, prompt_tokens.input_ids, end_tokens], dim=1) return tokenizer.decode(all_input_ids[0]) async def generate_speech(text, request_id): prompt_string = format_prompt(text) start = time.time() token_ids = [] async for output in engine.generate(prompt_string, sampling_params, request_id): token_ids = list(output.outputs[0].token_ids) gen_time = time.time() - start audio = decode_tokens_to_audio(token_ids, snac) if audio is None: return {'success': False, 'gen_time': gen_time} audio_duration = len(audio) / 24000 return { 'success': True, 'gen_time': gen_time, 'audio_duration': audio_duration, 'rtf': gen_time / audio_duration if audio_duration > 0 else float('inf') } test_texts = [ "Hello, how are you doing today?", "The weather is beautiful outside.", "I love programming with Python.", "Machine learning is fascinating.", "Can you help me with this task?", "Let me explain how this works.", "This is a test of the system.", "Technology is amazing these days.", "Have a wonderful day ahead.", "Thank you for your patience.", "Let's work together on this.", "The future looks very bright.", "I appreciate your help today.", "This demonstration is working.", "Audio generation is fast now.", "Real-time speech synthesis.", ] print("\n[2] Iniciando testes de stress...") print("=" * 70) results_summary = [] for num_users in [4, 8, 12, 16]: print(f"\n>>> TESTANDO {num_users} USUARIOS SIMULTANEOS <<<") print("-" * 50) texts = test_texts[:num_users] start_batch = time.time() tasks = [generate_speech(text, f"user_{i}_{num_users}") for i, text in enumerate(texts)] results = await asyncio.gather(*tasks) batch_time = time.time() - start_batch successful = [r for r in results if r['success']] total_audio = sum(r['audio_duration'] for r in successful) if len(successful) > 0: batch_rtf = batch_time / total_audio throughput = len(successful) / batch_time realtime_users = total_audio / batch_time print(f" Sucesso: {len(successful)}/{num_users}") print(f" Tempo total (wall-clock): {batch_time:.2f}s") print(f" Audio total gerado: {total_audio:.2f}s") print(f" Batch RTF: {batch_rtf:.3f}") print(f" Throughput: {throughput:.2f} req/s") print(f" USUARIOS REAL-TIME: {realtime_users:.1f}") results_summary.append({ 'users': num_users, 'success': len(successful), 'batch_time': batch_time, 'total_audio': total_audio, 'batch_rtf': batch_rtf, 'realtime_users': realtime_users }) else: print(f" ERRO: Nenhum audio gerado!") print("\n" + "=" * 70) print("RESUMO FINAL - CAPACIDADE RTX 4090") print("=" * 70) print("\n| Requests | Sucesso | Wall-Time | Audio Total | RTF | Real-time Users |") print("|----------|---------|-----------|-------------|-------|-----------------|") for r in results_summary: print(f"| {r['users']:8} | {r['success']:7} | {r['batch_time']:9.2f}s | {r['total_audio']:11.2f}s | {r['batch_rtf']:.3f} | {r['realtime_users']:15.1f} |") if results_summary: best = max(results_summary, key=lambda x: x['realtime_users']) print(f"\nMELHOR RESULTADO: {best['realtime_users']:.1f} usuarios real-time com {best['users']} requests") print("=" * 70) if __name__ == "__main__": asyncio.run(main())