OnyxMunk/AudioForge

1 2import asyncio3import httpx4import websockets5import json6import sys7 8API_URL = "http://localhost:8001/api/v1"9WS_URL = "ws://localhost:8001/api/v1"10 11async def demo_generation():12    print("🎵 AudioForge Real-Time Generation Demo 🎵")13    print("------------------------------------------")14 15    # 1. Create Generation Request16    prompt = "A fast-paced techno track with a thumping bassline"17    print(f"\n[1] Sending POST request...")18    print(f"    Prompt: '{prompt}'")19    20    async with httpx.AsyncClient() as client:21        try:22            response = await client.post(23                f"{API_URL}/generations/",24                json={"prompt": prompt, "duration": 5}  # Short duration for demo25            )26            response.raise_for_status()27            data = response.json()28            gen_id = data["id"]29            print(f"    ✅ Request Accepted! Generation ID: {gen_id}")30        except Exception as e:31            print(f"    ❌ Failed to create generation: {e}")32            return33 34    # 2. Connect to WebSocket35    print(f"\n[2] Connecting to WebSocket for updates...")36    ws_endpoint = f"{WS_URL}/ws/generations/{gen_id}"37    38    try:39        async with websockets.connect(ws_endpoint) as websocket:40            print("    ✅ Connected! Waiting for real-time updates...\n")41            print("    [STATUS]      [PROGRESS]    [MESSAGE]")42            print("    -------------------------------------")43            44            async for message in websocket:45                data = json.loads(message)46                47                # Format output nicely48                status = data.get("status", "unknown").upper()49                progress = f"{data.get('progress', 0)}%"50                msg = data.get("message", "")51                52                print(f"    {status:<13} {progress:<13} {msg}")53                54                if data.get("status") in ["completed", "failed"]:55                    if data.get("status") == "completed":56                        print(f"\n    🎉 Generation Complete! Audio URL: {data.get('audio_url')}")57                    else:58                        print(f"\n    ❌ Generation Failed: {data.get('error')}")59                    break60                    61    except Exception as e:62        print(f"\n    ❌ WebSocket Error: {e}")63 64if __name__ == "__main__":65    try:66        asyncio.run(demo_generation())67    except KeyboardInterrupt:68        print("\nDemo cancelled.")69