run formatter

This commit is contained in:
slSeanWU
2025-10-30 17:21:32 -04:00
parent 782ec8cbac
commit fdda8e98e0
4 changed files with 64 additions and 70 deletions
+31 -32
View File
@@ -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()