run formatter
This commit is contained in:
+31
-32
@@ -222,7 +222,7 @@ Examples:
|
||||
--prompts_file prompts.txt
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
# Required arguments
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
@@ -230,7 +230,7 @@ Examples:
|
||||
default="slseanwu/MIDI-LLM_Llama-3.2-1B",
|
||||
help="Path to MIDI-LLM model checkpoint, can be HuggingFace model ID or local path (default: slseanwu/MIDI-LLM_Llama-3.2-1B)"
|
||||
)
|
||||
|
||||
|
||||
# Input arguments (not required if using --interactive only)
|
||||
input_group = parser.add_mutually_exclusive_group(required=False)
|
||||
input_group.add_argument(
|
||||
@@ -243,7 +243,7 @@ Examples:
|
||||
type=str,
|
||||
help="Path to file containing prompts (one per line)"
|
||||
)
|
||||
|
||||
|
||||
# Output arguments
|
||||
parser.add_argument(
|
||||
"--output_root",
|
||||
@@ -257,7 +257,7 @@ Examples:
|
||||
default=DEFAULT_N_OUTPUTS,
|
||||
help=f"Number of outputs to generate per prompt (default: {DEFAULT_N_OUTPUTS})"
|
||||
)
|
||||
|
||||
|
||||
# Synthesis arguments
|
||||
parser.add_argument(
|
||||
"--no-synthesize",
|
||||
@@ -272,7 +272,7 @@ Examples:
|
||||
default="./soundfonts/FluidR3_GM/FluidR3_GM.sf2",
|
||||
help="Path to SoundFont file for synthesis (default: ./soundfonts/FluidR3_GM/FluidR3_GM.sf2)"
|
||||
)
|
||||
|
||||
|
||||
# Generation parameters
|
||||
parser.add_argument(
|
||||
"--temperature",
|
||||
@@ -292,7 +292,7 @@ Examples:
|
||||
default=DEFAULT_MAX_TOKENS,
|
||||
help=f"Maximum tokens to generate (default: {DEFAULT_MAX_TOKENS})"
|
||||
)
|
||||
|
||||
|
||||
# Model arguments
|
||||
parser.add_argument(
|
||||
"--cache_dir",
|
||||
@@ -305,13 +305,13 @@ Examples:
|
||||
action="store_true",
|
||||
help="Enter interactive mode after initial generation (keep generating until empty prompt)"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
# Validate that either prompts are provided or interactive mode is enabled
|
||||
if not args.prompt and not args.prompts_file and not args.interactive:
|
||||
parser.error("Either --prompt, --prompts_file, or --interactive must be specified")
|
||||
|
||||
|
||||
# Load prompts (if provided)
|
||||
prompts = []
|
||||
if args.prompt:
|
||||
@@ -320,7 +320,7 @@ Examples:
|
||||
with open(args.prompts_file, "r") as f:
|
||||
prompts = [line.strip() for line in f if line.strip()]
|
||||
print(f"Loaded {len(prompts)} prompts from {args.prompts_file}")
|
||||
|
||||
|
||||
# Check synthesis requirements
|
||||
if args.synthesize:
|
||||
soundfont_path = Path(args.soundfont)
|
||||
@@ -329,31 +329,31 @@ Examples:
|
||||
print("Please download a SoundFont or disable synthesis")
|
||||
import sys
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# Create output root directory with timestamp
|
||||
output_root = Path(args.output_root)
|
||||
session_timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
|
||||
output_dir = output_root / session_timestamp
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
print(f"Output directory: {output_dir.absolute()}\n")
|
||||
|
||||
|
||||
# Load tokenizer from the model checkpoint
|
||||
print("Loading tokenizer...")
|
||||
print("Loading tokenizer and model...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
args.model,
|
||||
cache_dir=args.cache_dir,
|
||||
pad_token="<|eot_id|>",
|
||||
)
|
||||
|
||||
|
||||
# Load model
|
||||
model = prepare_hf_model(model_path=args.model)
|
||||
|
||||
|
||||
# Generate from initial prompts (if provided)
|
||||
if prompts:
|
||||
print(f"Starting generation for {len(prompts)} prompt(s)...\n")
|
||||
start_time = time.time()
|
||||
|
||||
|
||||
stats = generate_from_prompts_hf(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
@@ -367,9 +367,9 @@ Examples:
|
||||
max_tokens=args.max_tokens,
|
||||
n_outputs=args.n_outputs
|
||||
)
|
||||
|
||||
|
||||
total_time = time.time() - start_time
|
||||
|
||||
|
||||
# Print summary
|
||||
print(f"\n{'='*70}")
|
||||
print("Generation Summary")
|
||||
@@ -378,22 +378,22 @@ Examples:
|
||||
print(f"Successful generations: {stats['successful_generations']}")
|
||||
print(f"Failed generations: {stats['failed_generations']}")
|
||||
print(f"Total time: {total_time:.2f}s")
|
||||
|
||||
|
||||
if stats['generation_times']:
|
||||
avg_time = sum(stats['generation_times']) / len(stats['generation_times'])
|
||||
print(f"Average generation time: {avg_time:.2f}s (excluding warmup)")
|
||||
|
||||
|
||||
print(f"\nOutputs saved to: {output_dir.absolute()}")
|
||||
|
||||
|
||||
# Print generated files
|
||||
if stats['output_files']:
|
||||
print(f"\nGenerated files:")
|
||||
for file_path in stats['output_files']:
|
||||
file_type = "🎵 MIDI" if file_path.endswith('.mid') else "🎧 Audio"
|
||||
print(f" {file_type}: {file_path}")
|
||||
|
||||
|
||||
print(f"{'='*70}\n")
|
||||
|
||||
|
||||
# Save stats to JSON
|
||||
stats_file = output_dir / "generation_stats.json"
|
||||
with open(stats_file, "w") as f:
|
||||
@@ -411,7 +411,7 @@ Examples:
|
||||
}, f, indent=2)
|
||||
else:
|
||||
print(f"No initial prompts provided. Starting in interactive mode...\n")
|
||||
|
||||
|
||||
# Interactive mode
|
||||
if args.interactive:
|
||||
print(f"\n{'='*70}")
|
||||
@@ -419,17 +419,17 @@ Examples:
|
||||
print(f"{'='*70}")
|
||||
print("Enter prompts to generate more MIDI files.")
|
||||
print("Press Enter with empty prompt to exit.\n")
|
||||
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Get user input
|
||||
user_prompt = input("Prompt: ").strip()
|
||||
|
||||
|
||||
# Exit if empty
|
||||
if not user_prompt:
|
||||
print("\nExiting interactive mode. Goodbye!")
|
||||
break
|
||||
|
||||
|
||||
# Generate from the new prompt
|
||||
print()
|
||||
interactive_stats = generate_from_prompts_hf(
|
||||
@@ -448,19 +448,19 @@ Examples:
|
||||
|
||||
# Print input prompt
|
||||
print(f"Input prompt: {user_prompt}")
|
||||
|
||||
|
||||
# Print mini summary
|
||||
print(f"\n✓ Generated {interactive_stats['successful_generations']}/{args.n_outputs} outputs")
|
||||
if interactive_stats['generation_times']:
|
||||
print(f" Generation time: {interactive_stats['generation_times'][0]:.2f}s")
|
||||
|
||||
|
||||
# Print file paths
|
||||
if interactive_stats['output_files']:
|
||||
for file_path in interactive_stats['output_files']:
|
||||
file_type = "🎵" if file_path.endswith('.mid') else "🎧"
|
||||
print(f" {file_type} {file_path}")
|
||||
print()
|
||||
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nInterrupted. Exiting interactive mode.")
|
||||
break
|
||||
@@ -471,4 +471,3 @@ Examples:
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user