ACE-Step/Ace-Step-v1.5

1#!/usr/bin/env python32"""3Batch Generate Text2Music Examples using LM4Generates 50 examples and saves them to examples/text2music/5"""6import os7import json8import sys9from pathlib import Path10 11# Add project root to path12project_root = Path(__file__).parent13sys.path.insert(0, str(project_root))14 15from acestep.llm_inference import LLMHandler16from loguru import logger17from tqdm import tqdm18 19 20def generate_examples(num_examples=50, output_dir="examples/text2music", start_index=1):21    """22    Generate examples using LM and save to JSON files23    24    Args:25        num_examples: Number of examples to generate26        output_dir: Output directory for JSON files27        start_index: Starting index for example files28    """29    # Initialize LLM Handler30    logger.info("Initializing LLM Handler...")31    llm_handler = LLMHandler()32    33    # Initialize LM34    checkpoint_dir = os.path.join(project_root, "checkpoints")35    36    # Use default LM model37    available_models = llm_handler.get_available_5hz_lm_models()38    if not available_models:39        logger.error("No 5Hz LM models found in checkpoints directory")40        return41    42    # Prefer acestep-5Hz-lm-0.6B if available43    lm_model = "acestep-5Hz-lm-0.6B" if "acestep-5Hz-lm-0.6B" in available_models else available_models[0]44    logger.info(f"Using LM model: {lm_model}")45    46    # Initialize LM47    status_msg, success = llm_handler.initialize(48        checkpoint_dir=checkpoint_dir,49        lm_model_path=lm_model,50        backend="vllm",  # Use vllm for faster generation51        device="auto",52        offload_to_cpu=False,53        dtype=None,54    )55    56    if not success:57        logger.error(f"Failed to initialize LM: {status_msg}")58        return59    60    logger.info(f"LM initialized successfully: {status_msg}")61    62    # Create output directory if it doesn't exist63    os.makedirs(output_dir, exist_ok=True)64    65    # Generate examples66    successful_count = 067    failed_count = 068    69    for i in tqdm(range(num_examples), desc="Generating examples"):70        example_num = start_index + i71        output_file = os.path.join(output_dir, f"example_{example_num:02d}.json")72        73        logger.info(f"Generating example {example_num}/{start_index + num_examples - 1}...")74        75        try:76            # Generate example using LM77            metadata, status = llm_handler.understand_audio_from_codes(78                audio_codes="NO USER INPUT",  # Empty input triggers example generation79                use_constrained_decoding=True,80                temperature=0.85,81                cfg_scale=1.0,82                top_k=None,83                top_p=0.9,84            )85            86            if not metadata:87                logger.warning(f"Failed to generate example {example_num}: {status}")88                failed_count += 189                continue90            91            # Build JSON data with all available fields92            example_data = {93                "think": True,  # Always true for LM-generated examples94                "caption": metadata.get("caption", ""),95                "lyrics": metadata.get("lyrics", ""),96            }97            98            # Add optional metadata fields if they exist and are not "N/A"99            if "bpm" in metadata and metadata["bpm"] not in [None, "N/A", ""]:100                try:101                    # Convert to int if it's a valid number102                    example_data["bpm"] = int(metadata["bpm"]) if isinstance(metadata["bpm"], (int, str)) else metadata["bpm"]103                except (ValueError, TypeError):104                    example_data["bpm"] = metadata["bpm"]105            106            if "duration" in metadata and metadata["duration"] not in [None, "N/A", ""]:107                try:108                    # Convert to int if it's a valid number109                    example_data["duration"] = int(metadata["duration"]) if isinstance(metadata["duration"], (int, str)) else metadata["duration"]110                except (ValueError, TypeError):111                    example_data["duration"] = metadata["duration"]112            113            if "keyscale" in metadata and metadata["keyscale"] not in [None, "N/A", ""]:114                example_data["keyscale"] = metadata["keyscale"]115            116            if "language" in metadata and metadata["language"] not in [None, "N/A", ""]:117                example_data["language"] = metadata["language"]118            119            if "timesignature" in metadata and metadata["timesignature"] not in [None, "N/A", ""]:120                example_data["timesignature"] = metadata["timesignature"]121            122            # Save to JSON file123            with open(output_file, 'w', encoding='utf-8') as f:124                json.dump(example_data, f, ensure_ascii=False, indent=4)125            126            logger.info(f"✅ Saved example {example_num} to {output_file}")127            logger.info(f"   Caption preview: {example_data['caption'][:100]}...")128            successful_count += 1129            130        except Exception as e:131            logger.error(f"❌ Error generating example {example_num}: {str(e)}")132            failed_count += 1133            continue134    135    # Summary136    logger.info(f"\n{'='*60}")137    logger.info(f"Generation complete!")138    logger.info(f"Successful: {successful_count}/{num_examples}")139    logger.info(f"Failed: {failed_count}/{num_examples}")140    logger.info(f"Output directory: {output_dir}")141    logger.info(f"{'='*60}\n")142 143 144if __name__ == "__main__":145    import argparse146    147    parser = argparse.ArgumentParser(description="Generate text2music examples using LM")148    parser.add_argument("--num", type=int, default=100, help="Number of examples to generate (default: 100)")149    parser.add_argument("--output-dir", type=str, default="examples/text2music", help="Output directory (default: examples/text2music)")150    parser.add_argument("--start-index", type=int, default=1, help="Starting index for example files (default: 1)")151    152    args = parser.parse_args()153    154    generate_examples(155        num_examples=args.num,156        output_dir=args.output_dir,157        start_index=args.start_index158    )159